github.com/annwntech/go-micro/v2@v2.9.5/tunnel/transport/transport.go (about)

     1  // Package transport provides a tunnel transport
     2  package transport
     3  
     4  import (
     5  	"context"
     6  
     7  	"github.com/annwntech/go-micro/v2/transport"
     8  	"github.com/annwntech/go-micro/v2/tunnel"
     9  )
    10  
    11  type tunTransport struct {
    12  	options transport.Options
    13  
    14  	tunnel tunnel.Tunnel
    15  }
    16  
    17  type tunnelKey struct{}
    18  
    19  type transportKey struct{}
    20  
    21  func (t *tunTransport) Init(opts ...transport.Option) error {
    22  	for _, o := range opts {
    23  		o(&t.options)
    24  	}
    25  
    26  	// close the existing tunnel
    27  	if t.tunnel != nil {
    28  		t.tunnel.Close()
    29  	}
    30  
    31  	// get the tunnel
    32  	tun, ok := t.options.Context.Value(tunnelKey{}).(tunnel.Tunnel)
    33  	if !ok {
    34  		tun = tunnel.NewTunnel()
    35  	}
    36  
    37  	// get the transport
    38  	tr, ok := t.options.Context.Value(transportKey{}).(transport.Transport)
    39  	if ok {
    40  		tun.Init(tunnel.Transport(tr))
    41  	}
    42  
    43  	// set the tunnel
    44  	t.tunnel = tun
    45  
    46  	return nil
    47  }
    48  
    49  func (t *tunTransport) Dial(addr string, opts ...transport.DialOption) (transport.Client, error) {
    50  	if err := t.tunnel.Connect(); err != nil {
    51  		return nil, err
    52  	}
    53  
    54  	c, err := t.tunnel.Dial(addr)
    55  	if err != nil {
    56  		return nil, err
    57  	}
    58  
    59  	return c, nil
    60  }
    61  
    62  func (t *tunTransport) Listen(addr string, opts ...transport.ListenOption) (transport.Listener, error) {
    63  	if err := t.tunnel.Connect(); err != nil {
    64  		return nil, err
    65  	}
    66  
    67  	l, err := t.tunnel.Listen(addr)
    68  	if err != nil {
    69  		return nil, err
    70  	}
    71  
    72  	return &tunListener{l}, nil
    73  }
    74  
    75  func (t *tunTransport) Options() transport.Options {
    76  	return t.options
    77  }
    78  
    79  func (t *tunTransport) String() string {
    80  	return "tunnel"
    81  }
    82  
    83  // NewTransport honours the initialiser used in
    84  func NewTransport(opts ...transport.Option) transport.Transport {
    85  	t := &tunTransport{
    86  		options: transport.Options{},
    87  	}
    88  
    89  	// initialise
    90  	t.Init(opts...)
    91  
    92  	return t
    93  }
    94  
    95  // WithTransport sets the internal tunnel
    96  func WithTunnel(t tunnel.Tunnel) transport.Option {
    97  	return func(o *transport.Options) {
    98  		if o.Context == nil {
    99  			o.Context = context.Background()
   100  		}
   101  		o.Context = context.WithValue(o.Context, tunnelKey{}, t)
   102  	}
   103  }
   104  
   105  // WithTransport sets the internal transport
   106  func WithTransport(t transport.Transport) transport.Option {
   107  	return func(o *transport.Options) {
   108  		if o.Context == nil {
   109  			o.Context = context.Background()
   110  		}
   111  		o.Context = context.WithValue(o.Context, transportKey{}, t)
   112  	}
   113  }