update context functions

proxy/context.go

@@ -10,9 +10,9 @@ type key int
 
 const (
 	sourceKey key = iota
-	destinationKey
-	originalDestinationKey
-	inboundDestinationKey
+	targetKey
+	originalTargetKey
+	inboundEntryPointKey
 	inboundTagKey
 	resolvedIPsKey
 )
@@ -21,60 +21,45 @@ func ContextWithSource(ctx context.Context, src net.Destination) context.Context
 	return context.WithValue(ctx, sourceKey, src)
 }
 
-func SourceFromContext(ctx context.Context) net.Destination {
-	v := ctx.Value(sourceKey)
-	if v == nil {
-		return net.Destination{}
-	}
-	return v.(net.Destination)
+func SourceFromContext(ctx context.Context) (net.Destination, bool) {
+	v, ok := ctx.Value(sourceKey).(net.Destination)
+	return v, ok
 }
 
-func ContextWithOriginalDestination(ctx context.Context, dest net.Destination) context.Context {
-	return context.WithValue(ctx, originalDestinationKey, dest)
+func ContextWithOriginalTarget(ctx context.Context, dest net.Destination) context.Context {
+	return context.WithValue(ctx, originalTargetKey, dest)
 }
 
-func OriginalDestinationFromContext(ctx context.Context) net.Destination {
-	v := ctx.Value(originalDestinationKey)
-	if v == nil {
-		return net.Destination{}
-	}
-	return v.(net.Destination)
+func OriginalTargetFromContext(ctx context.Context) (net.Destination, bool) {
+	v, ok := ctx.Value(originalTargetKey).(net.Destination)
+	return v, ok
 }
 
-func ContextWithDestination(ctx context.Context, dest net.Destination) context.Context {
-	return context.WithValue(ctx, destinationKey, dest)
+func ContextWithTarget(ctx context.Context, dest net.Destination) context.Context {
+	return context.WithValue(ctx, targetKey, dest)
 }
 
-func DestinationFromContext(ctx context.Context) net.Destination {
-	v := ctx.Value(destinationKey)
-	if v == nil {
-		return net.Destination{}
-	}
-	return v.(net.Destination)
+func TargetFromContext(ctx context.Context) (net.Destination, bool) {
+	v, ok := ctx.Value(targetKey).(net.Destination)
+	return v, ok
 }
 
-func ContextWithInboundDestination(ctx context.Context, dest net.Destination) context.Context {
-	return context.WithValue(ctx, inboundDestinationKey, dest)
+func ContextWithInboundEntryPoint(ctx context.Context, dest net.Destination) context.Context {
+	return context.WithValue(ctx, inboundEntryPointKey, dest)
 }
 
-func InboundDestinationFromContext(ctx context.Context) net.Destination {
-	v := ctx.Value(inboundDestinationKey)
-	if v == nil {
-		return net.Destination{}
-	}
-	return v.(net.Destination)
+func InboundEntryPointFromContext(ctx context.Context) (net.Destination, bool) {
+	v, ok := ctx.Value(inboundEntryPointKey).(net.Destination)
+	return v, ok
 }
 
 func ContextWithInboundTag(ctx context.Context, tag string) context.Context {
 	return context.WithValue(ctx, inboundTagKey, tag)
 }
 
-func InboundTagFromContext(ctx context.Context) string {
-	v := ctx.Value(inboundTagKey)
-	if v == nil {
-		return ""
-	}
-	return v.(string)
+func InboundTagFromContext(ctx context.Context) (string, bool) {
+	v, ok := ctx.Value(inboundTagKey).(string)
+	return v, ok
 }
 
 func ContextWithResolveIPs(ctx context.Context, ips []net.Address) context.Context {

proxy/dokodemo/dokodemo.go

@@ -52,7 +52,7 @@ func (d *DokodemoDoor) Process(ctx context.Context, network net.Network, conn in
 		Port:    d.port,
 	}
 	if d.config.FollowRedirect {
-		if origDest := proxy.OriginalDestinationFromContext(ctx); origDest.IsValid() {
+		if origDest, ok := proxy.OriginalTargetFromContext(ctx); ok {
 			dest = origDest
 		}
 	}

proxy/freedom/freedom.go

@@ -74,7 +74,7 @@ func (v *Handler) ResolveIP(destination net.Destination) net.Destination {
 }
 
 func (v *Handler) Process(ctx context.Context, outboundRay ray.OutboundRay, dialer proxy.Dialer) error {
-	destination := proxy.DestinationFromContext(ctx)
+	destination, _ := proxy.TargetFromContext(ctx)
 	if v.destOverride != nil {
 		server := v.destOverride.Server
 		destination = net.Destination{

proxy/shadowsocks/client.go

@@ -2,10 +2,9 @@ package shadowsocks
 
 import (
 	"context"
-
-	"time"
-
+	"errors"
 	"runtime"
+	"time"
 
 	"v2ray.com/core/app/log"
 	"v2ray.com/core/common"
@@ -40,7 +39,10 @@ func NewClient(ctx context.Context, config *ClientConfig) (*Client, error) {
 
 // Process implements OutboundHandler.Process().
 func (v *Client) Process(ctx context.Context, outboundRay ray.OutboundRay, dialer proxy.Dialer) error {
-	destination := proxy.DestinationFromContext(ctx)
+	destination, ok := proxy.TargetFromContext(ctx)
+	if !ok {
+		return errors.New("Shadowsocks|Client: Target not specified.")
+	}
 	network := destination.Network
 
 	var server *protocol.ServerSpec

proxy/shadowsocks/server.go

@@ -75,7 +75,6 @@ func (s *Server) Process(ctx context.Context, network net.Network, conn internet
 
 func (v *Server) handlerUDPPayload(ctx context.Context, conn internet.Connection, dispatcher dispatcher.Interface) error {
 	udpServer := udp.NewDispatcher(dispatcher)
-	source := proxy.SourceFromContext(ctx)
 
 	reader := buf.NewReader(conn)
 	for {
@@ -86,8 +85,10 @@ func (v *Server) handlerUDPPayload(ctx context.Context, conn internet.Connection
 
 		request, data, err := DecodeUDPPacket(v.user, payload)
 		if err != nil {
-			log.Info("Shadowsocks|Server: Skipping invalid UDP packet from: ", source, ": ", err)
-			log.Access(source, "", log.AccessRejected, err)
+			if source, ok := proxy.SourceFromContext(ctx); ok {
+				log.Info("Shadowsocks|Server: Skipping invalid UDP packet from: ", source, ": ", err)
+				log.Access(source, "", log.AccessRejected, err)
+			}
 			payload.Release()
 			continue
 		}
@@ -105,7 +106,9 @@ func (v *Server) handlerUDPPayload(ctx context.Context, conn internet.Connection
 		}
 
 		dest := request.Destination()
-		log.Access(source, dest, log.AccessAccepted, "")
+		if source, ok := proxy.SourceFromContext(ctx); ok {
+			log.Access(source, dest, log.AccessAccepted, "")
+		}
 		log.Info("Shadowsocks|Server: Tunnelling request to ", dest)
 
 		ctx = protocol.ContextWithUser(ctx, request.User)

proxy/socks/client.go

@@ -2,7 +2,7 @@ package socks
 
 import (
 	"context"
-
+	"errors"
 	"runtime"
 	"time"
 
@@ -35,7 +35,10 @@ func NewClient(ctx context.Context, config *ClientConfig) (*Client, error) {
 }
 
 func (c *Client) Process(ctx context.Context, ray ray.OutboundRay, dialer proxy.Dialer) error {
-	destination := proxy.DestinationFromContext(ctx)
+	destination, ok := proxy.TargetFromContext(ctx)
+	if !ok {
+		return errors.New("Socks|Client: Target not specified.")
+	}
 
 	var server *protocol.ServerSpec
 	var conn internet.Connection

proxy/socks/server.go

@@ -65,16 +65,20 @@ func (s *Server) processTCP(ctx context.Context, conn internet.Connection, dispa
 	conn.SetReadDeadline(time.Now().Add(time.Second * 8))
 	reader := bufio.NewReader(conn)
 
-	inboundDest := proxy.InboundDestinationFromContext(ctx)
+	inboundDest, ok := proxy.InboundEntryPointFromContext(ctx)
+	if !ok {
+		return errors.New("Socks|Server: inbound entry point not specified.")
+	}
 	session := &ServerSession{
 		config: s.config,
 		port:   inboundDest.Port,
 	}
 
-	source := proxy.SourceFromContext(ctx)
 	request, err := session.Handshake(reader, conn)
 	if err != nil {
-		log.Access(source, "", log.AccessRejected, err)
+		if source, ok := proxy.SourceFromContext(ctx); ok {
+			log.Access(source, "", log.AccessRejected, err)
+		}
 		log.Info("Socks|Server: Failed to read request: ", err)
 		return err
 	}
@@ -83,7 +87,9 @@ func (s *Server) processTCP(ctx context.Context, conn internet.Connection, dispa
 	if request.Command == protocol.RequestCommandTCP {
 		dest := request.Destination()
 		log.Info("Socks|Server: TCP Connect request to ", dest)
-		log.Access(source, dest, log.AccessAccepted, "")
+		if source, ok := proxy.SourceFromContext(ctx); ok {
+			log.Access(source, dest, log.AccessAccepted, "")
+		}
 
 		return s.transport(ctx, reader, conn, dest, dispatcher)
 	}
@@ -155,8 +161,9 @@ func (v *Server) transport(ctx context.Context, reader io.Reader, writer io.Writ
 func (v *Server) handleUDPPayload(ctx context.Context, conn internet.Connection, dispatcher dispatcher.Interface) error {
 	udpServer := udp.NewDispatcher(dispatcher)
 
-	source := proxy.SourceFromContext(ctx)
-	log.Info("Socks|Server: Client UDP connection from ", source)
+	if source, ok := proxy.SourceFromContext(ctx); ok {
+		log.Info("Socks|Server: Client UDP connection from ", source)
+	}
 
 	reader := buf.NewReader(conn)
 	for {
@@ -176,7 +183,9 @@ func (v *Server) handleUDPPayload(ctx context.Context, conn internet.Connection,
 		}
 
 		log.Info("Socks: Send packet to ", request.Destination(), " with ", len(data), " bytes")
-		log.Access(source, request.Destination, log.AccessAccepted, "")
+		if source, ok := proxy.SourceFromContext(ctx); ok {
+			log.Access(source, request.Destination, log.AccessAccepted, "")
+		}
 
 		dataBuf := buf.NewSmall()
 		dataBuf.Append(data)

proxy/vmess/outbound/outbound.go

@@ -67,7 +67,10 @@ func (v *VMessOutboundHandler) Process(ctx context.Context, outboundRay ray.Outb
 	}
 	defer conn.Close()
 
-	target := proxy.DestinationFromContext(ctx)
+	target, ok := proxy.TargetFromContext(ctx)
+	if !ok {
+		return errors.New("VMess|Outbound: Target not specified.")
+	}
 	log.Info("VMess|Outbound: Tunneling request to ", target, " via ", rec.Destination())
 
 	command := protocol.RequestCommandTCP