55 lines
1.0 KiB
Go
55 lines
1.0 KiB
Go
package comms
|
|
|
|
import (
|
|
"git.wamblee.org/converge/pkg/support/websocketutil"
|
|
"net"
|
|
)
|
|
|
|
type AgentListener struct {
|
|
decorated net.Listener
|
|
}
|
|
|
|
func NewAgentListener(listener net.Listener) AgentListener {
|
|
return AgentListener{decorated: listener}
|
|
}
|
|
|
|
type LocalAddrHackConn struct {
|
|
net.Conn
|
|
localAddr net.Addr
|
|
}
|
|
|
|
func (conn LocalAddrHackConn) LocalAddr() net.Addr {
|
|
return conn.localAddr
|
|
}
|
|
|
|
func NewLocalAddrHackConn(conn net.Conn, clientId string) LocalAddrHackConn {
|
|
addr := LocalAddrHackConn{
|
|
localAddr: websocketutil.WebSocketAddr(clientId),
|
|
}
|
|
addr.Conn = conn
|
|
return addr
|
|
}
|
|
|
|
func (listener AgentListener) Accept() (net.Conn, error) {
|
|
conn, err := listener.decorated.Accept()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
clientId, err := ReceiveClientInfo(conn)
|
|
if err != nil {
|
|
conn.Close()
|
|
return nil, err
|
|
}
|
|
conn = NewLocalAddrHackConn(conn, clientId.ClientId)
|
|
return conn, nil
|
|
}
|
|
|
|
func (listener AgentListener) Close() error {
|
|
return listener.decorated.Close()
|
|
}
|
|
|
|
func (listener AgentListener) Addr() net.Addr {
|
|
return listener.decorated.Addr()
|
|
}
|