118 lines
2.2 KiB
Go
118 lines
2.2 KiB
Go
package client
|
|
|
|
import (
|
|
"io"
|
|
"log"
|
|
"net"
|
|
pb "net-tunnel/pkg/proto"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type ConnectionState struct {
|
|
conn net.Conn
|
|
connID string
|
|
config *pb.ProxyConfig
|
|
stream pb.TunnelService_EstablishControlConnectionClient
|
|
closed bool
|
|
closedMutex sync.Mutex
|
|
closeOnce sync.Once
|
|
}
|
|
|
|
func NewConnectionState(conn net.Conn, connID string, config *pb.ProxyConfig, stream pb.TunnelService_EstablishControlConnectionClient) *ConnectionState {
|
|
return &ConnectionState{
|
|
conn: conn,
|
|
connID: connID,
|
|
config: config,
|
|
stream: stream,
|
|
closed: false,
|
|
closedMutex: sync.Mutex{},
|
|
closeOnce: sync.Once{},
|
|
}
|
|
}
|
|
|
|
func (c *ConnectionState) Close() {
|
|
c.closeOnce.Do(func() {
|
|
c.closedMutex.Lock()
|
|
defer c.closedMutex.Unlock()
|
|
|
|
if !c.closed {
|
|
c.closed = true
|
|
c.conn.Close()
|
|
log.Printf("Connection %s closed", c.connID)
|
|
}
|
|
})
|
|
}
|
|
|
|
func (c *ConnectionState) IsClosed() bool {
|
|
c.closedMutex.Lock()
|
|
defer c.closedMutex.Unlock()
|
|
return c.closed
|
|
}
|
|
|
|
func (c *ConnectionState) WriteData(data []byte) error {
|
|
c.closedMutex.Lock()
|
|
defer c.closedMutex.Unlock()
|
|
|
|
if c.closed {
|
|
return io.EOF
|
|
}
|
|
|
|
_, err := c.conn.Write(data)
|
|
if err != nil {
|
|
c.Close()
|
|
return err
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (c *ConnectionState) StartReading() {
|
|
go func() {
|
|
buffer := make([]byte, 4096)
|
|
defer c.Close()
|
|
|
|
for {
|
|
if c.IsClosed() {
|
|
return
|
|
}
|
|
|
|
if c.config.Type == pb.ProxyType_TCP {
|
|
if tcpConn, ok := c.conn.(*net.TCPConn); ok {
|
|
_ = tcpConn.SetReadDeadline(time.Now().Add(time.Minute * 3))
|
|
}
|
|
}
|
|
|
|
n, err := c.conn.Read(buffer)
|
|
if err != nil {
|
|
if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
|
|
continue
|
|
}
|
|
|
|
if err == io.EOF {
|
|
log.Printf("Connection %s closed by remote", c.connID)
|
|
}
|
|
return
|
|
}
|
|
|
|
if n > 0 {
|
|
if c.IsClosed() {
|
|
return
|
|
}
|
|
if err := c.stream.Send(&pb.Message{
|
|
Content: &pb.Message_ProxyData{
|
|
ProxyData: &pb.ProxyData{
|
|
ConnId: c.connID,
|
|
Data: buffer[:n],
|
|
ProxyConfig: c.config,
|
|
},
|
|
},
|
|
}); err != nil {
|
|
log.Printf("Failed to send proxy data: %v", err)
|
|
c.Close()
|
|
return
|
|
}
|
|
}
|
|
}
|
|
}()
|
|
}
|