From 41beb4ec04b84cef0acb628f9f5134d53c920205 Mon Sep 17 00:00:00 2001
From: RPRX <63339210+rprx@users.noreply.github.com>
Date: Sat, 5 Dec 2020 11:58:10 +0000
Subject: [PATCH] Add Timer to Inbound metadata struct

---
 common/session/session.go  | 5 ++++-
 proxy/dokodemo/dokodemo.go | 7 ++++++-
 proxy/http/server.go       | 8 ++++++--
 proxy/socks/server.go      | 8 ++++++--
 4 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/common/session/session.go b/common/session/session.go
index 683bf0b6..7b734897 100644
--- a/common/session/session.go
+++ b/common/session/session.go
@@ -8,6 +8,7 @@ import (
 	"github.com/xtls/xray-core/common/errors"
 	"github.com/xtls/xray-core/common/net"
 	"github.com/xtls/xray-core/common/protocol"
+	"github.com/xtls/xray-core/common/signal"
 )
 
 // ID of a session.
@@ -43,8 +44,10 @@ type Inbound struct {
 	Tag string
 	// User is the user that authencates for the inbound. May be nil if the protocol allows anounymous traffic.
 	User *protocol.MemoryUser
-	// Conn is actually internet.Connection.
+	// Conn is actually internet.Connection. May be nil.
 	Conn net.Conn
+	// Timer of the inbound buf copier. May be nil.
+	Timer *signal.ActivityTimer
 }
 
 // Outbound is the metadata of an outbound connection.
diff --git a/proxy/dokodemo/dokodemo.go b/proxy/dokodemo/dokodemo.go
index e9bb975f..8f7b70b9 100644
--- a/proxy/dokodemo/dokodemo.go
+++ b/proxy/dokodemo/dokodemo.go
@@ -103,7 +103,8 @@ func (d *DokodemoDoor) Process(ctx context.Context, network net.Network, conn in
 		return newError("unable to get destination")
 	}
 
-	if inbound := session.InboundFromContext(ctx); inbound != nil {
+	inbound := session.InboundFromContext(ctx)
+	if inbound != nil {
 		inbound.User = &protocol.MemoryUser{
 			Level: d.config.UserLevel,
 		}
@@ -121,6 +122,10 @@ func (d *DokodemoDoor) Process(ctx context.Context, network net.Network, conn in
 	ctx, cancel := context.WithCancel(ctx)
 	timer := signal.CancelAfterInactivity(ctx, cancel, plcy.Timeouts.ConnectionIdle)
 
+	if inbound != nil {
+		inbound.Timer = timer
+	}
+
 	ctx = policy.ContextWithBufferPolicy(ctx, plcy.Buffer)
 	link, err := dispatcher.Dispatch(ctx, dest)
 	if err != nil {
diff --git a/proxy/http/server.go b/proxy/http/server.go
index e3ff8d68..f7e37b2d 100644
--- a/proxy/http/server.go
+++ b/proxy/http/server.go
@@ -143,7 +143,7 @@ Start:
 	})
 
 	if strings.EqualFold(request.Method, "CONNECT") {
-		return s.handleConnect(ctx, request, reader, conn, dest, dispatcher)
+		return s.handleConnect(ctx, request, reader, conn, dest, dispatcher, inbound)
 	}
 
 	keepAlive := (strings.TrimSpace(strings.ToLower(request.Header.Get("Proxy-Connection"))) == "keep-alive")
@@ -159,7 +159,7 @@ Start:
 	return err
 }
 
-func (s *Server) handleConnect(ctx context.Context, _ *http.Request, reader *bufio.Reader, conn internet.Connection, dest net.Destination, dispatcher routing.Dispatcher) error {
+func (s *Server) handleConnect(ctx context.Context, _ *http.Request, reader *bufio.Reader, conn internet.Connection, dest net.Destination, dispatcher routing.Dispatcher, inbound *session.Inbound) error {
 	_, err := conn.Write([]byte("HTTP/1.1 200 Connection established\r\n\r\n"))
 	if err != nil {
 		return newError("failed to write back OK response").Base(err)
@@ -169,6 +169,10 @@ func (s *Server) handleConnect(ctx context.Context, _ *http.Request, reader *buf
 	ctx, cancel := context.WithCancel(ctx)
 	timer := signal.CancelAfterInactivity(ctx, cancel, plcy.Timeouts.ConnectionIdle)
 
+	if inbound != nil {
+		inbound.Timer = timer
+	}
+
 	ctx = policy.ContextWithBufferPolicy(ctx, plcy.Buffer)
 	link, err := dispatcher.Dispatch(ctx, dest)
 	if err != nil {
diff --git a/proxy/socks/server.go b/proxy/socks/server.go
index de24ad1e..f947b3d1 100644
--- a/proxy/socks/server.go
+++ b/proxy/socks/server.go
@@ -128,7 +128,7 @@ func (s *Server) processTCP(ctx context.Context, conn internet.Connection, dispa
 			})
 		}
 
-		return s.transport(ctx, reader, conn, dest, dispatcher)
+		return s.transport(ctx, reader, conn, dest, dispatcher, inbound)
 	}
 
 	if request.Command == protocol.RequestCommandUDP {
@@ -144,10 +144,14 @@ func (*Server) handleUDP(c io.Reader) error {
 	return common.Error2(io.Copy(buf.DiscardBytes, c))
 }
 
-func (s *Server) transport(ctx context.Context, reader io.Reader, writer io.Writer, dest net.Destination, dispatcher routing.Dispatcher) error {
+func (s *Server) transport(ctx context.Context, reader io.Reader, writer io.Writer, dest net.Destination, dispatcher routing.Dispatcher, inbound *session.Inbound) error {
 	ctx, cancel := context.WithCancel(ctx)
 	timer := signal.CancelAfterInactivity(ctx, cancel, s.policy().Timeouts.ConnectionIdle)
 
+	if inbound != nil {
+		inbound.Timer = timer
+	}
+
 	plcy := s.policy()
 	ctx = policy.ContextWithBufferPolicy(ctx, plcy.Buffer)
 	link, err := dispatcher.Dispatch(ctx, dest)