github.com/xmplusdev/xmcore@v1.8.11-0.20240412132628-5518b55526af/common/mux/server.go (about)

     1  package mux
     2  
     3  import (
     4  	"context"
     5  	"io"
     6  
     7  	"github.com/xmplusdev/xmcore/common"
     8  	"github.com/xmplusdev/xmcore/common/buf"
     9  	"github.com/xmplusdev/xmcore/common/errors"
    10  	"github.com/xmplusdev/xmcore/common/log"
    11  	"github.com/xmplusdev/xmcore/common/net"
    12  	"github.com/xmplusdev/xmcore/common/protocol"
    13  	"github.com/xmplusdev/xmcore/common/session"
    14  	"github.com/xmplusdev/xmcore/core"
    15  	"github.com/xmplusdev/xmcore/features/routing"
    16  	"github.com/xmplusdev/xmcore/transport"
    17  	"github.com/xmplusdev/xmcore/transport/pipe"
    18  )
    19  
    20  type Server struct {
    21  	dispatcher routing.Dispatcher
    22  }
    23  
    24  // NewServer creates a new mux.Server.
    25  func NewServer(ctx context.Context) *Server {
    26  	s := &Server{}
    27  	core.RequireFeatures(ctx, func(d routing.Dispatcher) {
    28  		s.dispatcher = d
    29  	})
    30  	return s
    31  }
    32  
    33  // Type implements common.HasType.
    34  func (s *Server) Type() interface{} {
    35  	return s.dispatcher.Type()
    36  }
    37  
    38  // Dispatch implements routing.Dispatcher
    39  func (s *Server) Dispatch(ctx context.Context, dest net.Destination) (*transport.Link, error) {
    40  	if dest.Address != muxCoolAddress {
    41  		return s.dispatcher.Dispatch(ctx, dest)
    42  	}
    43  
    44  	opts := pipe.OptionsFromContext(ctx)
    45  	uplinkReader, uplinkWriter := pipe.New(opts...)
    46  	downlinkReader, downlinkWriter := pipe.New(opts...)
    47  
    48  	_, err := NewServerWorker(ctx, s.dispatcher, &transport.Link{
    49  		Reader: uplinkReader,
    50  		Writer: downlinkWriter,
    51  	})
    52  	if err != nil {
    53  		return nil, err
    54  	}
    55  
    56  	return &transport.Link{Reader: downlinkReader, Writer: uplinkWriter}, nil
    57  }
    58  
    59  // DispatchLink implements routing.Dispatcher
    60  func (s *Server) DispatchLink(ctx context.Context, dest net.Destination, link *transport.Link) error {
    61  	if dest.Address != muxCoolAddress {
    62  		return s.dispatcher.DispatchLink(ctx, dest, link)
    63  	}
    64  	_, err := NewServerWorker(ctx, s.dispatcher, link)
    65  	return err
    66  }
    67  
    68  // Start implements common.Runnable.
    69  func (s *Server) Start() error {
    70  	return nil
    71  }
    72  
    73  // Close implements common.Closable.
    74  func (s *Server) Close() error {
    75  	return nil
    76  }
    77  
    78  type ServerWorker struct {
    79  	dispatcher     routing.Dispatcher
    80  	link           *transport.Link
    81  	sessionManager *SessionManager
    82  }
    83  
    84  func NewServerWorker(ctx context.Context, d routing.Dispatcher, link *transport.Link) (*ServerWorker, error) {
    85  	worker := &ServerWorker{
    86  		dispatcher:     d,
    87  		link:           link,
    88  		sessionManager: NewSessionManager(),
    89  	}
    90  	go worker.run(ctx)
    91  	return worker, nil
    92  }
    93  
    94  func handle(ctx context.Context, s *Session, output buf.Writer) {
    95  	writer := NewResponseWriter(s.ID, output, s.transferType)
    96  	if err := buf.Copy(s.input, writer); err != nil {
    97  		newError("session ", s.ID, " ends.").Base(err).WriteToLog(session.ExportIDToError(ctx))
    98  		writer.hasError = true
    99  	}
   100  
   101  	writer.Close()
   102  	s.Close(false)
   103  }
   104  
   105  func (w *ServerWorker) ActiveConnections() uint32 {
   106  	return uint32(w.sessionManager.Size())
   107  }
   108  
   109  func (w *ServerWorker) Closed() bool {
   110  	return w.sessionManager.Closed()
   111  }
   112  
   113  func (w *ServerWorker) handleStatusKeepAlive(meta *FrameMetadata, reader *buf.BufferedReader) error {
   114  	if meta.Option.Has(OptionData) {
   115  		return buf.Copy(NewStreamReader(reader), buf.Discard)
   116  	}
   117  	return nil
   118  }
   119  
   120  func (w *ServerWorker) handleStatusNew(ctx context.Context, meta *FrameMetadata, reader *buf.BufferedReader) error {
   121  	newError("received request for ", meta.Target).WriteToLog(session.ExportIDToError(ctx))
   122  	{
   123  		msg := &log.AccessMessage{
   124  			To:     meta.Target,
   125  			Status: log.AccessAccepted,
   126  			Reason: "",
   127  		}
   128  		if inbound := session.InboundFromContext(ctx); inbound != nil && inbound.Source.IsValid() {
   129  			msg.From = inbound.Source
   130  			msg.Email = inbound.User.Email
   131  		}
   132  		ctx = log.ContextWithAccessMessage(ctx, msg)
   133  	}
   134  
   135  	if network := session.AllowedNetworkFromContext(ctx); network != net.Network_Unknown {
   136  		if meta.Target.Network != network {
   137  			return newError("unexpected network ", meta.Target.Network) // it will break the whole Mux connection
   138  		}
   139  	}
   140  
   141  	if meta.GlobalID != [8]byte{} { // MUST ignore empty Global ID
   142  		mb, err := NewPacketReader(reader, &meta.Target).ReadMultiBuffer()
   143  		if err != nil {
   144  			return err
   145  		}
   146  		XUDPManager.Lock()
   147  		x := XUDPManager.Map[meta.GlobalID]
   148  		if x == nil {
   149  			x = &XUDP{GlobalID: meta.GlobalID}
   150  			XUDPManager.Map[meta.GlobalID] = x
   151  			XUDPManager.Unlock()
   152  		} else {
   153  			if x.Status == Initializing { // nearly impossible
   154  				XUDPManager.Unlock()
   155  				newError("XUDP hit ", meta.GlobalID).Base(errors.New("conflict")).AtWarning().WriteToLog(session.ExportIDToError(ctx))
   156  				// It's not a good idea to return an err here, so just let client wait.
   157  				// Client will receive an End frame after sending a Keep frame.
   158  				return nil
   159  			}
   160  			x.Status = Initializing
   161  			XUDPManager.Unlock()
   162  			x.Mux.Close(false) // detach from previous Mux
   163  			b := buf.New()
   164  			b.Write(mb[0].Bytes())
   165  			b.UDP = mb[0].UDP
   166  			if err = x.Mux.output.WriteMultiBuffer(mb); err != nil {
   167  				x.Interrupt()
   168  				mb = buf.MultiBuffer{b}
   169  			} else {
   170  				b.Release()
   171  				mb = nil
   172  			}
   173  			newError("XUDP hit ", meta.GlobalID).Base(err).WriteToLog(session.ExportIDToError(ctx))
   174  		}
   175  		if mb != nil {
   176  			ctx = session.ContextWithTimeoutOnly(ctx, true)
   177  			// Actually, it won't return an error in Xray-core's implementations.
   178  			link, err := w.dispatcher.Dispatch(ctx, meta.Target)
   179  			if err != nil {
   180  				XUDPManager.Lock()
   181  				delete(XUDPManager.Map, x.GlobalID)
   182  				XUDPManager.Unlock()
   183  				err = newError("XUDP new ", meta.GlobalID).Base(errors.New("failed to dispatch request to ", meta.Target).Base(err))
   184  				return err // it will break the whole Mux connection
   185  			}
   186  			link.Writer.WriteMultiBuffer(mb) // it's meaningless to test a new pipe
   187  			x.Mux = &Session{
   188  				input:  link.Reader,
   189  				output: link.Writer,
   190  			}
   191  			newError("XUDP new ", meta.GlobalID).Base(err).WriteToLog(session.ExportIDToError(ctx))
   192  		}
   193  		x.Mux = &Session{
   194  			input:        x.Mux.input,
   195  			output:       x.Mux.output,
   196  			parent:       w.sessionManager,
   197  			ID:           meta.SessionID,
   198  			transferType: protocol.TransferTypePacket,
   199  			XUDP:         x,
   200  		}
   201  		go handle(ctx, x.Mux, w.link.Writer)
   202  		x.Status = Active
   203  		if !w.sessionManager.Add(x.Mux) {
   204  			x.Mux.Close(false)
   205  		}
   206  		return nil
   207  	}
   208  
   209  	link, err := w.dispatcher.Dispatch(ctx, meta.Target)
   210  	if err != nil {
   211  		if meta.Option.Has(OptionData) {
   212  			buf.Copy(NewStreamReader(reader), buf.Discard)
   213  		}
   214  		return newError("failed to dispatch request.").Base(err)
   215  	}
   216  	s := &Session{
   217  		input:        link.Reader,
   218  		output:       link.Writer,
   219  		parent:       w.sessionManager,
   220  		ID:           meta.SessionID,
   221  		transferType: protocol.TransferTypeStream,
   222  	}
   223  	if meta.Target.Network == net.Network_UDP {
   224  		s.transferType = protocol.TransferTypePacket
   225  	}
   226  	w.sessionManager.Add(s)
   227  	go handle(ctx, s, w.link.Writer)
   228  	if !meta.Option.Has(OptionData) {
   229  		return nil
   230  	}
   231  
   232  	rr := s.NewReader(reader, &meta.Target)
   233  	if err := buf.Copy(rr, s.output); err != nil {
   234  		buf.Copy(rr, buf.Discard)
   235  		return s.Close(false)
   236  	}
   237  	return nil
   238  }
   239  
   240  func (w *ServerWorker) handleStatusKeep(meta *FrameMetadata, reader *buf.BufferedReader) error {
   241  	if !meta.Option.Has(OptionData) {
   242  		return nil
   243  	}
   244  
   245  	s, found := w.sessionManager.Get(meta.SessionID)
   246  	if !found {
   247  		// Notify remote peer to close this session.
   248  		closingWriter := NewResponseWriter(meta.SessionID, w.link.Writer, protocol.TransferTypeStream)
   249  		closingWriter.Close()
   250  
   251  		return buf.Copy(NewStreamReader(reader), buf.Discard)
   252  	}
   253  
   254  	rr := s.NewReader(reader, &meta.Target)
   255  	err := buf.Copy(rr, s.output)
   256  
   257  	if err != nil && buf.IsWriteError(err) {
   258  		newError("failed to write to downstream writer. closing session ", s.ID).Base(err).WriteToLog()
   259  		s.Close(false)
   260  		return buf.Copy(rr, buf.Discard)
   261  	}
   262  
   263  	return err
   264  }
   265  
   266  func (w *ServerWorker) handleStatusEnd(meta *FrameMetadata, reader *buf.BufferedReader) error {
   267  	if s, found := w.sessionManager.Get(meta.SessionID); found {
   268  		s.Close(false)
   269  	}
   270  	if meta.Option.Has(OptionData) {
   271  		return buf.Copy(NewStreamReader(reader), buf.Discard)
   272  	}
   273  	return nil
   274  }
   275  
   276  func (w *ServerWorker) handleFrame(ctx context.Context, reader *buf.BufferedReader) error {
   277  	var meta FrameMetadata
   278  	err := meta.Unmarshal(reader)
   279  	if err != nil {
   280  		return newError("failed to read metadata").Base(err)
   281  	}
   282  
   283  	switch meta.SessionStatus {
   284  	case SessionStatusKeepAlive:
   285  		err = w.handleStatusKeepAlive(&meta, reader)
   286  	case SessionStatusEnd:
   287  		err = w.handleStatusEnd(&meta, reader)
   288  	case SessionStatusNew:
   289  		err = w.handleStatusNew(ctx, &meta, reader)
   290  	case SessionStatusKeep:
   291  		err = w.handleStatusKeep(&meta, reader)
   292  	default:
   293  		status := meta.SessionStatus
   294  		return newError("unknown status: ", status).AtError()
   295  	}
   296  
   297  	if err != nil {
   298  		return newError("failed to process data").Base(err)
   299  	}
   300  	return nil
   301  }
   302  
   303  func (w *ServerWorker) run(ctx context.Context) {
   304  	input := w.link.Reader
   305  	reader := &buf.BufferedReader{Reader: input}
   306  
   307  	defer w.sessionManager.Close()
   308  
   309  	for {
   310  		select {
   311  		case <-ctx.Done():
   312  			return
   313  		default:
   314  			err := w.handleFrame(ctx, reader)
   315  			if err != nil {
   316  				if errors.Cause(err) != io.EOF {
   317  					newError("unexpected EOF").Base(err).WriteToLog(session.ExportIDToError(ctx))
   318  					common.Interrupt(input)
   319  				}
   320  				return
   321  			}
   322  		}
   323  	}
   324  }