improvements

Novage · Feb 26, 2024 · 22bdd17 · 22bdd17
1 parent 3d7fb37
commit 22bdd17
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 22 deletions.
diff --git a/packages/p2p-media-loader-core/src/p2p/peer-protocol.ts b/packages/p2p-media-loader-core/src/p2p/peer-protocol.ts
@@ -1,14 +1,14 @@
 import { PeerConnection } from "bittorrent-tracker";
-import * as Command from "./commands";
-import * as Utils from "../utils/utils";
 import { Settings } from "../types";
+import * as Utils from "../utils/utils";
+import * as Command from "./commands";
 
 export type PeerSettings = Pick<
   Settings,
   | "p2pNotReceivingBytesTimeoutMs"
   | "webRtcMaxMessageSize"
   | "p2pErrorRetries"
-  | "validateP2Psegment"
+  | "validateP2PSegment"
 >;
 
 export class PeerProtocol {

diff --git a/packages/p2p-media-loader-core/src/p2p/peer.ts b/packages/p2p-media-loader-core/src/p2p/peer.ts
@@ -1,16 +1,16 @@
 import { PeerConnection } from "bittorrent-tracker";
-import { PeerProtocol, PeerSettings } from "./peer-protocol";
+import debug from "debug";
 import {
+  PeerRequestErrorType,
   Request,
   RequestControls,
   RequestError,
-  PeerRequestErrorType,
   RequestInnerErrorType,
 } from "../requests/request";
-import * as Command from "./commands";
 import { Segment } from "../types";
 import * as Utils from "../utils/utils";
-import debug from "debug";
+import * as Command from "./commands";
+import { PeerProtocol, PeerSettings } from "./peer-protocol";
 
 const { PeerCommandType } = Command;
 type PeerEventHandlers = {
@@ -86,7 +86,9 @@ export class Peer {
             request.setTotalBytes(command.s);
           } else if (request.totalBytes - request.loadedBytes !== command.s) {
             request.clearLoadedBytes();
-            this.cancelSegmentDownloading("peer-response-bytes-mismatch");
+            this.cancelSegmentDownloading(
+              "peer-response-bytes-length-mismatch",
+            );
             this.destroy();
           }
         }
@@ -107,17 +109,24 @@ export class Peer {
           return;
         }
 
-        let isValid = true;
-        if (this.settings.validateP2Psegment) {
-          isValid = await this.settings.validateP2Psegment(
+        const isWrongBytes = request.loadedBytes !== request.totalBytes;
+
+        if (isWrongBytes) {
+          request.clearLoadedBytes();
+          this.cancelSegmentDownloading("peer-response-bytes-length-mismatch");
+          this.destroy();
+          return;
+        }
+
+        const isValid =
+          (await this.settings.validateP2PSegment?.(
             request.segment.url,
             request.segment.byteRange,
-          );
-        }
+          )) ?? true;
 
         if (!isValid) {
           request.clearLoadedBytes();
-          this.cancelSegmentDownloading("peer-response-bytes-mismatch");
+          this.cancelSegmentDownloading("p2p-segment-validation-failed");
           this.destroy();
           return;
         }
@@ -152,7 +161,7 @@ export class Peer {
 
     if (isOverflow) {
       request.clearLoadedBytes();
-      this.cancelSegmentDownloading("peer-response-bytes-mismatch");
+      this.cancelSegmentDownloading("peer-response-bytes-length-mismatch");
       this.destroy();
       return;
     }

diff --git a/packages/p2p-media-loader-core/src/requests/request.ts b/packages/p2p-media-loader-core/src/requests/request.ts
@@ -1,8 +1,8 @@
-import { Segment, Playback, BandwidthCalculators } from "../types";
+import debug from "debug";
+import { BandwidthCalculators, Playback, Segment } from "../types";
+import * as LoggerUtils from "../utils/logger";
 import * as StreamUtils from "../utils/stream";
 import * as Utils from "../utils/utils";
-import * as LoggerUtils from "../utils/logger";
-import debug from "debug";
 
 export type LoadProgress = {
   startTimestamp: number;
@@ -323,10 +323,11 @@ export type HttpRequestErrorType =
   | "http-unexpected-status-code";
 
 export type PeerRequestErrorType =
-  | "peer-response-bytes-mismatch"
+  | "peer-response-bytes-length-mismatch"
   | "peer-protocol-violation"
   | "peer-segment-absent"
-  | "peer-closed";
+  | "peer-closed"
+  | "p2p-segment-validation-failed";
 
 type RequestErrorType =
   | RequestInnerErrorType

diff --git a/packages/p2p-media-loader-core/src/types.d.ts b/packages/p2p-media-loader-core/src/types.d.ts
@@ -1,5 +1,5 @@
-import { RequestAttempt } from "./requests/request";
 import { BandwidthCalculator } from "./bandwidth-calculator";
+import { RequestAttempt } from "./requests/request";
 
 export type StreamType = "main" | "secondary";
 
@@ -53,7 +53,7 @@ export type Settings = {
   httpNotReceivingBytesTimeoutMs: number;
   httpErrorRetries: number;
   p2pErrorRetries: number;
-  validateP2Psegment?: (url: string, byteRange?: ByteRange) => Promise<boolean>;
+  validateP2PSegment?: (url: string, byteRange?: ByteRange) => Promise<boolean>;
 };
 
 export type CoreEventHandlers = {