From 88b25d38cbd9dd40bbb47dde863d1a4ed7118889 Mon Sep 17 00:00:00 2001 From: Darien Raymond Date: Wed, 21 Feb 2018 17:05:29 +0100 Subject: [PATCH] simplify context retrieval --- app/commander/commander.go | 5 +---- app/dispatcher/default.go | 6 +----- app/dns/server.go | 6 +----- app/log/command/command.go | 5 +---- app/proxyman/command/command.go | 5 +---- app/proxyman/inbound/dynamic.go | 5 +---- app/proxyman/inbound/inbound.go | 5 +---- app/proxyman/mux/mux.go | 2 +- app/proxyman/outbound/handler.go | 5 +---- app/proxyman/outbound/outbound.go | 5 +---- app/router/router.go | 6 +----- context.go | 11 ++++++++++- proxy/dokodemo/dokodemo.go | 6 +----- proxy/freedom/freedom.go | 6 +----- proxy/http/server.go | 5 +---- proxy/shadowsocks/client.go | 6 +----- proxy/shadowsocks/server.go | 6 +----- proxy/socks/client.go | 6 +----- proxy/socks/server.go | 5 +---- proxy/vmess/inbound/inbound.go | 6 +----- proxy/vmess/outbound/outbound.go | 6 +----- 21 files changed, 30 insertions(+), 88 deletions(-) diff --git a/app/commander/commander.go b/app/commander/commander.go index e1e408a97..19991fce4 100644 --- a/app/commander/commander.go +++ b/app/commander/commander.go @@ -22,10 +22,7 @@ type Commander struct { } func NewCommander(ctx context.Context, config *Config) (*Commander, error) { - v := core.FromContext(ctx) - if v == nil { - return nil, newError("V is not in context.") - } + v := core.MustFromContext(ctx) c := &Commander{ config: *config, ohm: v.OutboundHandlerManager(), diff --git a/app/dispatcher/default.go b/app/dispatcher/default.go index 6be84fcf3..1aba8be66 100644 --- a/app/dispatcher/default.go +++ b/app/dispatcher/default.go @@ -27,11 +27,7 @@ type DefaultDispatcher struct { // NewDefaultDispatcher create a new DefaultDispatcher. func NewDefaultDispatcher(ctx context.Context, config *Config) (*DefaultDispatcher, error) { - v := core.FromContext(ctx) - if v == nil { - return nil, newError("V is not in context.") - } - + v := core.MustFromContext(ctx) d := &DefaultDispatcher{ ohm: v.OutboundHandlerManager(), router: v.Router(), diff --git a/app/dns/server.go b/app/dns/server.go index 2972ff5cf..866c5ff0f 100644 --- a/app/dns/server.go +++ b/app/dns/server.go @@ -54,11 +54,7 @@ func New(ctx context.Context, config *Config) (*Server, error) { return nil }, } - v := core.FromContext(ctx) - if v == nil { - return nil, newError("V is not in context.") - } - + v := core.MustFromContext(ctx) if err := v.RegisterFeature((*core.DNSClient)(nil), server); err != nil { return nil, newError("unable to register DNSClient.").Base(err) } diff --git a/app/log/command/command.go b/app/log/command/command.go index 84d3078c9..7e7804996 100644 --- a/app/log/command/command.go +++ b/app/log/command/command.go @@ -41,10 +41,7 @@ func (s *service) Register(server *grpc.Server) { func init() { common.Must(common.RegisterConfig((*Config)(nil), func(ctx context.Context, cfg interface{}) (interface{}, error) { - s := core.FromContext(ctx) - if s == nil { - return nil, newError("V is not in context.") - } + s := core.MustFromContext(ctx) return &service{v: s}, nil })) } diff --git a/app/proxyman/command/command.go b/app/proxyman/command/command.go index c0357579d..023b83c5e 100644 --- a/app/proxyman/command/command.go +++ b/app/proxyman/command/command.go @@ -139,10 +139,7 @@ func (s *service) Register(server *grpc.Server) { func init() { common.Must(common.RegisterConfig((*Config)(nil), func(ctx context.Context, cfg interface{}) (interface{}, error) { - s := core.FromContext(ctx) - if s == nil { - return nil, newError("V is not in context.") - } + s := core.MustFromContext(ctx) return &service{v: s}, nil })) } diff --git a/app/proxyman/inbound/dynamic.go b/app/proxyman/inbound/dynamic.go index e6f77972b..ca2cc9440 100644 --- a/app/proxyman/inbound/dynamic.go +++ b/app/proxyman/inbound/dynamic.go @@ -29,10 +29,7 @@ type DynamicInboundHandler struct { } func NewDynamicInboundHandler(ctx context.Context, tag string, receiverConfig *proxyman.ReceiverConfig, proxyConfig interface{}) (*DynamicInboundHandler, error) { - v := core.FromContext(ctx) - if v == nil { - return nil, newError("V is not in context.") - } + v := core.MustFromContext(ctx) h := &DynamicInboundHandler{ tag: tag, proxyConfig: proxyConfig, diff --git a/app/proxyman/inbound/inbound.go b/app/proxyman/inbound/inbound.go index e7a190bb4..c752216f9 100644 --- a/app/proxyman/inbound/inbound.go +++ b/app/proxyman/inbound/inbound.go @@ -24,10 +24,7 @@ func New(ctx context.Context, config *proxyman.InboundConfig) (*Manager, error) m := &Manager{ taggedHandlers: make(map[string]core.InboundHandler), } - v := core.FromContext(ctx) - if v == nil { - return nil, newError("V is not in context") - } + v := core.MustFromContext(ctx) if err := v.RegisterFeature((*core.InboundHandlerManager)(nil), m); err != nil { return nil, newError("unable to register InboundHandlerManager").Base(err) } diff --git a/app/proxyman/mux/mux.go b/app/proxyman/mux/mux.go index b921dbd69..85a1ae5ac 100644 --- a/app/proxyman/mux/mux.go +++ b/app/proxyman/mux/mux.go @@ -261,7 +261,7 @@ type Server struct { // NewServer creates a new mux.Server. func NewServer(ctx context.Context) *Server { s := &Server{ - dispatcher: core.FromContext(ctx).Dispatcher(), + dispatcher: core.MustFromContext(ctx).Dispatcher(), } return s } diff --git a/app/proxyman/outbound/handler.go b/app/proxyman/outbound/handler.go index 63a4ec55b..c2cf8ebdf 100644 --- a/app/proxyman/outbound/handler.go +++ b/app/proxyman/outbound/handler.go @@ -22,10 +22,7 @@ type Handler struct { } func NewHandler(ctx context.Context, config *core.OutboundHandlerConfig) (*Handler, error) { - v := core.FromContext(ctx) - if v == nil { - return nil, newError("V is not in context") - } + v := core.MustFromContext(ctx) h := &Handler{ config: config, outboundManager: v.OutboundHandlerManager(), diff --git a/app/proxyman/outbound/outbound.go b/app/proxyman/outbound/outbound.go index 52848f324..91bc2c441 100644 --- a/app/proxyman/outbound/outbound.go +++ b/app/proxyman/outbound/outbound.go @@ -25,10 +25,7 @@ func New(ctx context.Context, config *proxyman.OutboundConfig) (*Manager, error) m := &Manager{ taggedHandler: make(map[string]core.OutboundHandler), } - v := core.FromContext(ctx) - if v == nil { - return nil, newError("V is not in context") - } + v := core.MustFromContext(ctx) if err := v.RegisterFeature((*core.OutboundHandlerManager)(nil), m); err != nil { return nil, newError("unable to register OutboundHandlerManager").Base(err) } diff --git a/app/router/router.go b/app/router/router.go index 5e0109d5b..4267064c5 100644 --- a/app/router/router.go +++ b/app/router/router.go @@ -18,11 +18,7 @@ type Router struct { } func NewRouter(ctx context.Context, config *Config) (*Router, error) { - v := core.FromContext(ctx) - if v == nil { - return nil, newError("V is not in context") - } - + v := core.MustFromContext(ctx) r := &Router{ domainStrategy: config.DomainStrategy, rules: make([]Rule, len(config.Rule)), diff --git a/context.go b/context.go index c92f0bacc..2651f3fe2 100644 --- a/context.go +++ b/context.go @@ -8,10 +8,19 @@ type key int const v2rayKey key = 1 -// FromContext returns a Instance from the given context, or nil if the context doesn't contain one. +// FromContext returns an Instance from the given context, or nil if the context doesn't contain one. func FromContext(ctx context.Context) *Instance { if s, ok := ctx.Value(v2rayKey).(*Instance); ok { return s } return nil } + +// MustFromContext returns an Instance from the given context, or panics if not present. +func MustFromContext(ctx context.Context) *Instance { + v := FromContext(ctx) + if v == nil { + panic("V is not in context.") + } + return v +} diff --git a/proxy/dokodemo/dokodemo.go b/proxy/dokodemo/dokodemo.go index 8f834b9a2..78be47c12 100644 --- a/proxy/dokodemo/dokodemo.go +++ b/proxy/dokodemo/dokodemo.go @@ -28,11 +28,7 @@ func New(ctx context.Context, config *Config) (*DokodemoDoor, error) { if config.NetworkList == nil || config.NetworkList.Size() == 0 { return nil, newError("no network specified") } - v := core.FromContext(ctx) - if v == nil { - return nil, newError("V is not in context.") - } - + v := core.MustFromContext(ctx) d := &DokodemoDoor{ config: config, address: config.GetPredefinedAddress(), diff --git a/proxy/freedom/freedom.go b/proxy/freedom/freedom.go index 80292fb60..a45f0a70d 100644 --- a/proxy/freedom/freedom.go +++ b/proxy/freedom/freedom.go @@ -27,11 +27,7 @@ type Handler struct { // New creates a new Freedom handler. func New(ctx context.Context, config *Config) (*Handler, error) { - v := core.FromContext(ctx) - if v == nil { - return nil, newError("V is not found in context.") - } - + v := core.MustFromContext(ctx) f := &Handler{ config: *config, policyManager: v.PolicyManager(), diff --git a/proxy/http/server.go b/proxy/http/server.go index 9179de132..2acc426b8 100644 --- a/proxy/http/server.go +++ b/proxy/http/server.go @@ -31,10 +31,7 @@ type Server struct { func NewServer(ctx context.Context, config *ServerConfig) (*Server, error) { s := &Server{ config: config, - v: core.FromContext(ctx), - } - if s.v == nil { - return nil, newError("V is not in context.") + v: core.MustFromContext(ctx), } return s, nil diff --git a/proxy/shadowsocks/client.go b/proxy/shadowsocks/client.go index 58e659db3..317a6588d 100644 --- a/proxy/shadowsocks/client.go +++ b/proxy/shadowsocks/client.go @@ -32,12 +32,8 @@ func NewClient(ctx context.Context, config *ClientConfig) (*Client, error) { } client := &Client{ serverPicker: protocol.NewRoundRobinServerPicker(serverList), - v: core.FromContext(ctx), + v: core.MustFromContext(ctx), } - if client.v == nil { - return nil, newError("V is not in context.") - } - return client, nil } diff --git a/proxy/shadowsocks/server.go b/proxy/shadowsocks/server.go index 31b345a07..96fb65107 100644 --- a/proxy/shadowsocks/server.go +++ b/proxy/shadowsocks/server.go @@ -39,11 +39,7 @@ func NewServer(ctx context.Context, config *ServerConfig) (*Server, error) { config: config, user: config.GetUser(), account: account, - v: core.FromContext(ctx), - } - - if s.v == nil { - return nil, newError("V is not in context.") + v: core.MustFromContext(ctx), } return s, nil diff --git a/proxy/socks/client.go b/proxy/socks/client.go index 78191e1e6..61d6495fd 100644 --- a/proxy/socks/client.go +++ b/proxy/socks/client.go @@ -32,11 +32,7 @@ func NewClient(ctx context.Context, config *ClientConfig) (*Client, error) { return nil, newError("0 target server") } - v := core.FromContext(ctx) - if v == nil { - return nil, newError("V is not in context") - } - + v := core.MustFromContext(ctx) return &Client{ serverPicker: protocol.NewRoundRobinServerPicker(serverList), policyManager: v.PolicyManager(), diff --git a/proxy/socks/server.go b/proxy/socks/server.go index 262ef4ad5..5b0524edd 100644 --- a/proxy/socks/server.go +++ b/proxy/socks/server.go @@ -27,10 +27,7 @@ type Server struct { func NewServer(ctx context.Context, config *ServerConfig) (*Server, error) { s := &Server{ config: config, - v: core.FromContext(ctx), - } - if s.v == nil { - return nil, newError("V is not in context.") + v: core.MustFromContext(ctx), } return s, nil } diff --git a/proxy/vmess/inbound/inbound.go b/proxy/vmess/inbound/inbound.go index 145a60ba2..792bd3438 100644 --- a/proxy/vmess/inbound/inbound.go +++ b/proxy/vmess/inbound/inbound.go @@ -105,11 +105,7 @@ type Handler struct { // New creates a new VMess inbound handler. func New(ctx context.Context, config *Config) (*Handler, error) { - v := core.FromContext(ctx) - if v == nil { - return nil, newError("V is not in context.") - } - + v := core.MustFromContext(ctx) handler := &Handler{ policyManager: v.PolicyManager(), inboundHandlerManager: v.InboundHandlerManager(), diff --git a/proxy/vmess/outbound/outbound.go b/proxy/vmess/outbound/outbound.go index 3225cb7ac..e4cb0faa2 100644 --- a/proxy/vmess/outbound/outbound.go +++ b/proxy/vmess/outbound/outbound.go @@ -35,11 +35,7 @@ func New(ctx context.Context, config *Config) (*Handler, error) { handler := &Handler{ serverList: serverList, serverPicker: protocol.NewRoundRobinServerPicker(serverList), - v: core.FromContext(ctx), - } - - if handler.v == nil { - return nil, newError("V is not in context.") + v: core.MustFromContext(ctx), } return handler, nil