github.com/v2fly/v2ray-core/v5@v5.16.2-0.20240507031116-8191faa6e095/common/session/context.go (about)

     1  package session
     2  
     3  import (
     4  	"context"
     5  )
     6  
     7  type sessionKey int
     8  
     9  const (
    10  	idSessionKey sessionKey = iota
    11  	inboundSessionKey
    12  	outboundSessionKey
    13  	contentSessionKey
    14  	muxPreferedSessionKey
    15  	sockoptSessionKey
    16  	trackedConnectionErrorKey
    17  	handlerSessionKey // nolint: varcheck
    18  )
    19  
    20  // ContextWithID returns a new context with the given ID.
    21  func ContextWithID(ctx context.Context, id ID) context.Context {
    22  	return context.WithValue(ctx, idSessionKey, id)
    23  }
    24  
    25  // IDFromContext returns ID in this context, or 0 if not contained.
    26  func IDFromContext(ctx context.Context) ID {
    27  	if id, ok := ctx.Value(idSessionKey).(ID); ok {
    28  		return id
    29  	}
    30  	return 0
    31  }
    32  
    33  func ContextWithInbound(ctx context.Context, inbound *Inbound) context.Context {
    34  	return context.WithValue(ctx, inboundSessionKey, inbound)
    35  }
    36  
    37  func InboundFromContext(ctx context.Context) *Inbound {
    38  	if inbound, ok := ctx.Value(inboundSessionKey).(*Inbound); ok {
    39  		return inbound
    40  	}
    41  	return nil
    42  }
    43  
    44  func ContextWithOutbound(ctx context.Context, outbound *Outbound) context.Context {
    45  	return context.WithValue(ctx, outboundSessionKey, outbound)
    46  }
    47  
    48  func OutboundFromContext(ctx context.Context) *Outbound {
    49  	if outbound, ok := ctx.Value(outboundSessionKey).(*Outbound); ok {
    50  		return outbound
    51  	}
    52  	return nil
    53  }
    54  
    55  func ContextWithContent(ctx context.Context, content *Content) context.Context {
    56  	return context.WithValue(ctx, contentSessionKey, content)
    57  }
    58  
    59  func ContentFromContext(ctx context.Context) *Content {
    60  	if content, ok := ctx.Value(contentSessionKey).(*Content); ok {
    61  		return content
    62  	}
    63  	return nil
    64  }
    65  
    66  // ContextWithMuxPrefered returns a new context with the given bool
    67  func ContextWithMuxPrefered(ctx context.Context, forced bool) context.Context {
    68  	return context.WithValue(ctx, muxPreferedSessionKey, forced)
    69  }
    70  
    71  // MuxPreferedFromContext returns value in this context, or false if not contained.
    72  func MuxPreferedFromContext(ctx context.Context) bool {
    73  	if val, ok := ctx.Value(muxPreferedSessionKey).(bool); ok {
    74  		return val
    75  	}
    76  	return false
    77  }
    78  
    79  // ContextWithSockopt returns a new context with Socket configs included
    80  func ContextWithSockopt(ctx context.Context, s *Sockopt) context.Context {
    81  	return context.WithValue(ctx, sockoptSessionKey, s)
    82  }
    83  
    84  // SockoptFromContext returns Socket configs in this context, or nil if not contained.
    85  func SockoptFromContext(ctx context.Context) *Sockopt {
    86  	if sockopt, ok := ctx.Value(sockoptSessionKey).(*Sockopt); ok {
    87  		return sockopt
    88  	}
    89  	return nil
    90  }
    91  
    92  func GetTransportLayerProxyTagFromContext(ctx context.Context) string {
    93  	if ContentFromContext(ctx) == nil {
    94  		return ""
    95  	}
    96  	return ContentFromContext(ctx).Attribute("transportLayerOutgoingTag")
    97  }
    98  
    99  func SetTransportLayerProxyTagToContext(ctx context.Context, tag string) context.Context {
   100  	if contentFromContext := ContentFromContext(ctx); contentFromContext == nil {
   101  		ctx = ContextWithContent(ctx, &Content{})
   102  	}
   103  	ContentFromContext(ctx).SetAttribute("transportLayerOutgoingTag", tag)
   104  	return ctx
   105  }
   106  
   107  func GetForcedOutboundTagFromContext(ctx context.Context) string {
   108  	if ContentFromContext(ctx) == nil {
   109  		return ""
   110  	}
   111  	return ContentFromContext(ctx).Attribute("forcedOutboundTag")
   112  }
   113  
   114  func SetForcedOutboundTagToContext(ctx context.Context, tag string) context.Context {
   115  	if contentFromContext := ContentFromContext(ctx); contentFromContext == nil {
   116  		ctx = ContextWithContent(ctx, &Content{})
   117  	}
   118  	ContentFromContext(ctx).SetAttribute("forcedOutboundTag", tag)
   119  	return ctx
   120  }
   121  
   122  type TrackedRequestErrorFeedback interface {
   123  	SubmitError(err error)
   124  }
   125  
   126  func SubmitOutboundErrorToOriginator(ctx context.Context, err error) {
   127  	if errorTracker := ctx.Value(trackedConnectionErrorKey); errorTracker != nil {
   128  		errorTracker := errorTracker.(TrackedRequestErrorFeedback)
   129  		errorTracker.SubmitError(err)
   130  	}
   131  }
   132  
   133  func TrackedConnectionError(ctx context.Context, tracker TrackedRequestErrorFeedback) context.Context {
   134  	return context.WithValue(ctx, trackedConnectionErrorKey, tracker)
   135  }