世界 il y a 2 ans
Parent
commit
6dcacf3b5e
1 fichiers modifiés avec 21 ajouts et 15 suppressions
  1. 21 15
      common/sniff/sniff.go

+ 21 - 15
common/sniff/sniff.go

@@ -22,23 +22,29 @@ func PeekStream(ctx context.Context, conn net.Conn, buffer *buf.Buffer, timeout
 	if timeout == 0 {
 		timeout = C.ReadPayloadTimeout
 	}
-	err := conn.SetReadDeadline(time.Now().Add(timeout))
-	if err != nil {
-		return nil, E.Cause(err, "set read deadline")
-	}
-	_, err = buffer.ReadOnceFrom(conn)
-	err = E.Errors(err, conn.SetReadDeadline(time.Time{}))
-	if err != nil {
-		return nil, E.Cause(err, "read payload")
-	}
-	var metadata *adapter.InboundContext
+	deadline := time.Now().Add(timeout)
 	var errors []error
-	for _, sniffer := range sniffers {
-		metadata, err = sniffer(ctx, bytes.NewReader(buffer.Bytes()))
-		if metadata != nil {
-			return metadata, nil
+
+	for i := 0; i < 3; i++ {
+		err := conn.SetReadDeadline(deadline)
+		if err != nil {
+			return nil, E.Cause(err, "set read deadline")
+		}
+		_, err = buffer.ReadOnceFrom(conn)
+		err = E.Errors(err, conn.SetReadDeadline(time.Time{}))
+		if err != nil {
+			if i > 0 {
+				break
+			}
+			return nil, E.Cause(err, "read payload")
+		}
+		for _, sniffer := range sniffers {
+			metadata, err := sniffer(ctx, bytes.NewReader(buffer.Bytes()))
+			if metadata != nil {
+				return metadata, nil
+			}
+			errors = append(errors, err)
 		}
-		errors = append(errors, err)
 	}
 	return nil, E.Errors(errors...)
 }