feat: add a timeout option for tcp detector

mrhaoxx · Aug 13, 2024 · f59461d · f59461d
1 parent f704425
commit f59461d
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 3 deletions.
diff --git a/tcp/detect.go b/tcp/detect.go
@@ -7,6 +7,7 @@ import (
  "io"
  "net/http"
  "reflect"
+ "time"
 
  utils "github.com/mrhaoxx/OpenNG/utils"
 )
@@ -40,6 +41,9 @@ type Detector func(io.Reader, *Conn) string
 
 type Detect struct {
  Dets []Detector
+
+ Timeout time.Duration
+ TimeoutProtocol string
 }
 
 // var detectors = []detector{
@@ -65,11 +69,19 @@ func (det *Detect) Handle(c *Conn) SerRet {
  Writer: raw,
  Rawconn: raw,
  })
+ if det.Timeout > 0 {
+ raw.SetReadDeadline(time.Now().Add(det.Timeout))
+ defer raw.SetReadDeadline(time.Time{})
+ }
  var proto string
- for _, f := range det.Dets {
+ for i, f := range det.Dets {
  if proto = f(buf, c); proto != "" {
  break
  }
+ if i == 0 && buf.buffer.Len() == 0 {
+ proto = det.TimeoutProtocol
+ break
+ }
  buf.Reset(true)
  }
  if proto == "" {

diff --git a/ui/builtin.go b/ui/builtin.go
@@ -247,6 +247,14 @@ var _builtin_refs_assertions = map[string]Assert{
  },
  },
  },
+ "timeout": {
+ Type: "int",
+ Default: 0,
+ },
+ "timeoutprotocol": {
+ Type: "string",
+ Default: "UNKNOWN",
+ },
  },
  },
  "builtin::tcp::controller": {
@@ -879,6 +887,10 @@ var _builtin_refs = map[string]Inst{
  },
  "builtin::tcp::det": func(spec *ArgNode) (any, error) {
  protocols := spec.MustGet("protocols").ToStringList()
+ timeout := spec.MustGet("timeout").ToInt()
+ _t := time.Duration(timeout)
+
+ timeoutprotocol := spec.MustGet("timeoutprotocol").ToString()
 
  var dets []tcp.Detector
  for _, p := range protocols {
@@ -898,9 +910,9 @@ var _builtin_refs = map[string]Inst{
  }
  }
 
- log.Verboseln(fmt.Sprintf("new tcp detector: protocols=%#v", protocols))
+ log.Verboseln(fmt.Sprintf("new tcp detector: protocols=%#v timeout=%d(%s) timeoutprotocol=%s", protocols, timeout, _t.String(), timeoutprotocol))
 
- return &tcp.Detect{Dets: dets}, nil
+ return &tcp.Detect{Dets: dets, Timeout: _t, TimeoutProtocol: timeoutprotocol}, nil
  },
  "builtin::tcp::controller": func(spec *ArgNode) (any, error) {
  services := spec.MustGet("services").ToMap()