github.com/searKing/golang/go@v1.2.117/net/mux/conn.go (about)

     1  // Copyright 2020 The searKing Author. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package mux
     6  
     7  import (
     8  	"context"
     9  	"errors"
    10  	"net"
    11  	"runtime"
    12  	"sync"
    13  	"time"
    14  
    15  	"github.com/searKing/golang/go/sync/atomic"
    16  )
    17  
    18  var (
    19  	// ErrHijacked is returned by ResponseWriter.Write calls when
    20  	// the underlying connection has been hijacked using the
    21  	// Hijacker interface. A zero-byte write on a hijacked
    22  	// connection will return ErrHijacked without any other side
    23  	// effects.
    24  	ErrHijacked = errors.New("cmux: connection has been hijacked")
    25  )
    26  
    27  type conn struct {
    28  	// server is the server on which the connection arrived.
    29  	// Immutable; never nil.
    30  	server *Server
    31  
    32  	// cancelCtx cancels the connection-level context.
    33  	cancelCtx context.CancelFunc
    34  
    35  	// rwc is the underlying network connection.
    36  	muc *sniffConn
    37  
    38  	// remoteAddr is rwc.RemoteAddr().String(). It is not populated synchronously
    39  	// inside the Listener's Accept goroutine, as some implementations block.
    40  	// It is populated immediately inside the (*conn).serve goroutine.
    41  	// This is the value of a HandlerConn's (*Request).RemoteAddr.
    42  	remoteAddr string
    43  
    44  	curPacketState atomic.Uint64 // packed (unixtime<<8|uint8(ConnStateHook))
    45  
    46  	// mu guards hijackedv
    47  	mu sync.Mutex
    48  
    49  	// hijackedv is whether this connection has been hijacked
    50  	// by a HandlerConn with the Hijacker interface.
    51  	// It is guarded by mu.
    52  	hijackedv bool
    53  }
    54  
    55  func (c *conn) hijacked() bool {
    56  	c.mu.Lock()
    57  	defer c.mu.Unlock()
    58  	return c.hijackedv
    59  }
    60  
    61  // c.mu must be held.
    62  func (c *conn) hijackLocked() (rwc net.Conn, err error) {
    63  	if c.hijackedv {
    64  		return nil, ErrHijacked
    65  	}
    66  
    67  	c.hijackedv = true
    68  	rwc = c.muc
    69  
    70  	c.setState(rwc, ConnStateHijacked)
    71  	return
    72  }
    73  
    74  // Close the connection.
    75  func (c *conn) close() {
    76  	_ = c.muc.Close()
    77  }
    78  
    79  // Serve a new connection.
    80  func (c *conn) serve(ctx context.Context) {
    81  	c.remoteAddr = c.muc.RemoteAddr().String()
    82  	ctx = context.WithValue(ctx, LocalAddrContextKey, c.muc.LocalAddr())
    83  	defer func() {
    84  		if err := recover(); err != nil && err != ErrAbortHandler {
    85  			const size = 64 << 10
    86  			buf := make([]byte, size)
    87  			buf = buf[:runtime.Stack(buf, false)]
    88  			c.server.logf("mux: panic serving %v: %v\n%s", c.remoteAddr, err, buf)
    89  		}
    90  		if !c.hijacked() {
    91  			c.close()
    92  			c.setState(c.muc, ConnStateClosed)
    93  		}
    94  	}()
    95  
    96  	ctx, cancelCtx := context.WithCancel(ctx)
    97  	c.cancelCtx = cancelCtx
    98  	defer cancelCtx()
    99  
   100  	c.setState(c.muc, ConnStateActive)
   101  
   102  	rwc, err := c.hijackLocked() // so the conn is taken over
   103  	if err != nil {
   104  		return
   105  	}
   106  	serverHandler{c.server}.Serve(rwc)
   107  	c.setState(c.muc, ConnStateIdle)
   108  }
   109  
   110  func (c *conn) setState(nc net.Conn, state ConnState) {
   111  	srv := c.server
   112  	switch state {
   113  	case ConnStateNew:
   114  		srv.trackConn(c, true)
   115  	case ConnStateHijacked, ConnStateClosed:
   116  		srv.trackConn(c, false)
   117  	}
   118  	if state > 0xff || state < 0 {
   119  		panic("internal error")
   120  	}
   121  	packedState := uint64(time.Now().Unix()<<8) | uint64(state)
   122  	c.curPacketState.Store(packedState)
   123  	if hook := srv.ConnStateHook; hook != nil {
   124  		hook(nc, state)
   125  	}
   126  }
   127  
   128  func (c *conn) getState() (state ConnState, unixSec int64) {
   129  	packedState := c.curPacketState.Load()
   130  	return ConnState(packedState & 0xff), int64(packedState >> 8)
   131  }