customizable policy

proxy/socks/server.go

@@ -8,6 +8,7 @@ import (
 	"v2ray.com/core/app"
 	"v2ray.com/core/app/dispatcher"
 	"v2ray.com/core/app/log"
+	"v2ray.com/core/app/policy"
 	"v2ray.com/core/common"
 	"v2ray.com/core/common/buf"
 	"v2ray.com/core/common/net"
@@ -21,6 +22,7 @@ import (
 // Server is a SOCKS 5 proxy server
 type Server struct {
 	config *ServerConfig
+	policy policy.Policy
 }
 
 // NewServer creates a new Server object.
@@ -32,6 +34,17 @@ func NewServer(ctx context.Context, config *ServerConfig) (*Server, error) {
 	s := &Server{
 		config: config,
 	}
+	space.OnInitialize(func() error {
+		pm := policy.PolicyFromSpace(space)
+		if pm == nil {
+			return newError("Policy not found in space.")
+		}
+		s.policy = pm.GetPolicy(config.UserLevel)
+		if config.Timeout > 0 && config.UserLevel == 0 {
+			s.policy.Timeout.ConnectionIdle.Value = config.Timeout
+		}
+		return nil
+	})
 	return s, nil
 }
 
@@ -57,7 +70,7 @@ func (s *Server) Process(ctx context.Context, network net.Network, conn internet
 }
 
 func (s *Server) processTCP(ctx context.Context, conn internet.Connection, dispatcher dispatcher.Interface) error {
-	conn.SetReadDeadline(time.Now().Add(time.Second * 8))
+	conn.SetReadDeadline(time.Now().Add(s.policy.Timeout.Handshake.Duration()))
 	reader := buf.NewBufferedReader(buf.NewReader(conn))
 
 	inboundDest, ok := proxy.InboundEntryPointFromContext(ctx)
@@ -103,12 +116,8 @@ func (*Server) handleUDP(c net.Conn) error {
 }
 
 func (v *Server) transport(ctx context.Context, reader io.Reader, writer io.Writer, dest net.Destination, dispatcher dispatcher.Interface) error {
-	timeout := time.Second * time.Duration(v.config.Timeout)
-	if timeout == 0 {
-		timeout = time.Minute * 5
-	}
 	ctx, cancel := context.WithCancel(ctx)
-	timer := signal.CancelAfterInactivity(ctx, cancel, timeout)
+	timer := signal.CancelAfterInactivity(ctx, cancel, v.policy.Timeout.ConnectionIdle.Duration())
 
 	ray, err := dispatcher.Dispatch(ctx, dest)
 	if err != nil {
@@ -125,6 +134,7 @@ func (v *Server) transport(ctx context.Context, reader io.Reader, writer io.Writ
 		if err := buf.Copy(v2reader, input, buf.UpdateActivity(timer)); err != nil {
 			return newError("failed to transport all TCP request").Base(err)
 		}
+		timer.SetTimeout(v.policy.Timeout.DownlinkOnly.Duration())
 		return nil
 	})
 
@@ -133,7 +143,7 @@ func (v *Server) transport(ctx context.Context, reader io.Reader, writer io.Writ
 		if err := buf.Copy(output, v2writer, buf.UpdateActivity(timer)); err != nil {
 			return newError("failed to transport all TCP response").Base(err)
 		}
-		timer.SetTimeout(time.Second * 2)
+		timer.SetTimeout(v.policy.Timeout.UplinkOnly.Duration())
 		return nil
 	})