feat: initial commit
This commit is contained in:
redtea committed
1 parent
353f62e257
commit
c6852c6f1b
62 files changed
+5911
-2
No files matched your search
@@ -0,0 +1,374 @@
|
||||
package agent
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/binary"
|
||||
"fmt"
|
||||
"log"
|
||||
"net"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/x-grk/longpolling-remote-access/internal/tunnel"
|
||||
)
|
||||
|
||||
// --- direct TCP mux tunnel ---
|
||||
|
||||
// runTunnel keeps a persistent connection to the host's 8446 tunnel, dialing
|
||||
// client-side ports on demand for "tcp" forward rules.
|
||||
func (a *Agent) runTunnel(ctx context.Context) {
|
||||
backoff := time.Second
|
||||
for ctx.Err() == nil {
|
||||
if err := a.tunnelOnce(ctx); err != nil {
|
||||
select {
|
||||
case <-time.After(backoff):
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
if backoff < 15*time.Second {
|
||||
backoff *= 2
|
||||
}
|
||||
continue
|
||||
}
|
||||
backoff = time.Second
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Agent) tunnelOnce(ctx context.Context) error {
|
||||
d := net.Dialer{Timeout: 10 * time.Second}
|
||||
conn, err := d.DialContext(ctx, "tcp", a.cfg.TunnelAddr)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer conn.Close()
|
||||
if tc, ok := conn.(*net.TCPConn); ok {
|
||||
_ = tc.SetNoDelay(true)
|
||||
}
|
||||
|
||||
// Handshake: server sends "LPRA-TUN1 <nonceHex>", we reply "<agentID> <sig>".
|
||||
_ = conn.SetDeadline(time.Now().Add(10 * time.Second))
|
||||
line, err := readLineConn(conn, 512)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
f := strings.Fields(line)
|
||||
if len(f) != 2 || f[0] != "LPRA-TUN1" {
|
||||
return fmt.Errorf("bad tunnel greeting")
|
||||
}
|
||||
nonce, err := hexDecode(f[1])
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
key := tunnel.AuthKey(a.state.Token)
|
||||
if _, err := fmt.Fprintf(conn, "%s %s\n", a.state.AgentID, tunnel.Sign(key, nonce)); err != nil {
|
||||
return err
|
||||
}
|
||||
_ = conn.SetDeadline(time.Time{})
|
||||
log.Printf("tunnel: connected to %s", a.cfg.TunnelAddr)
|
||||
|
||||
tc := &tunClient{conn: conn, streams: map[uint32]*agentStream{}}
|
||||
return tc.readLoop()
|
||||
}
|
||||
|
||||
type tunClient struct {
|
||||
conn net.Conn
|
||||
wmu sync.Mutex
|
||||
mu sync.Mutex
|
||||
streams map[uint32]*agentStream
|
||||
}
|
||||
|
||||
// agentStream buffers downstream data in a channel so bytes that arrive before
|
||||
// the client-side dial completes are not lost (the OPEN/DATA race).
|
||||
type agentStream struct {
|
||||
in chan []byte
|
||||
closed chan struct{}
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (s *agentStream) close() { s.once.Do(func() { close(s.closed) }) }
|
||||
|
||||
func (t *tunClient) write(f tunnel.Frame) error {
|
||||
t.wmu.Lock()
|
||||
defer t.wmu.Unlock()
|
||||
return tunnel.WriteFrame(t.conn, f)
|
||||
}
|
||||
|
||||
func (t *tunClient) readLoop() error {
|
||||
for {
|
||||
f, err := tunnel.ReadFrame(t.conn)
|
||||
if err != nil {
|
||||
t.shutdown()
|
||||
return err
|
||||
}
|
||||
switch f.Type {
|
||||
case tunnel.TypeOpen:
|
||||
if len(f.Payload) >= 2 {
|
||||
port := int(f.Payload[0])<<8 | int(f.Payload[1])
|
||||
t.mu.Lock()
|
||||
s := &agentStream{in: make(chan []byte, 256), closed: make(chan struct{})}
|
||||
t.streams[f.Stream] = s
|
||||
t.mu.Unlock()
|
||||
go t.openStream(f.Stream, port, s)
|
||||
}
|
||||
case tunnel.TypeData:
|
||||
t.mu.Lock()
|
||||
s := t.streams[f.Stream]
|
||||
t.mu.Unlock()
|
||||
if s != nil {
|
||||
select {
|
||||
case s.in <- f.Payload:
|
||||
case <-s.closed:
|
||||
}
|
||||
}
|
||||
case tunnel.TypeClose:
|
||||
t.closeStream(f.Stream)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (t *tunClient) openStream(id uint32, port int, s *agentStream) {
|
||||
conn, err := net.DialTimeout("tcp", net.JoinHostPort("127.0.0.1", itoa(port)), 10*time.Second)
|
||||
if err != nil {
|
||||
_ = t.write(tunnel.Frame{Type: tunnel.TypeClose, Stream: id})
|
||||
t.closeStream(id)
|
||||
return
|
||||
}
|
||||
if tc, ok := conn.(*net.TCPConn); ok {
|
||||
_ = tc.SetNoDelay(true)
|
||||
}
|
||||
defer conn.Close()
|
||||
|
||||
// downstream buffer -> local socket
|
||||
go func() {
|
||||
for {
|
||||
select {
|
||||
case b := <-s.in:
|
||||
if _, err := conn.Write(b); err != nil {
|
||||
s.close()
|
||||
return
|
||||
}
|
||||
case <-s.closed:
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
// local socket -> upstream frames
|
||||
buf := make([]byte, 32*1024)
|
||||
for {
|
||||
n, err := conn.Read(buf)
|
||||
if n > 0 {
|
||||
if t.write(tunnel.Frame{Type: tunnel.TypeData, Stream: id, Payload: append([]byte(nil), buf[:n]...)}) != nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
_ = t.write(tunnel.Frame{Type: tunnel.TypeClose, Stream: id})
|
||||
t.closeStream(id)
|
||||
}
|
||||
|
||||
func (t *tunClient) closeStream(id uint32) {
|
||||
t.mu.Lock()
|
||||
s := t.streams[id]
|
||||
delete(t.streams, id)
|
||||
t.mu.Unlock()
|
||||
if s != nil {
|
||||
s.close()
|
||||
}
|
||||
}
|
||||
|
||||
func (t *tunClient) shutdown() {
|
||||
t.mu.Lock()
|
||||
streams := t.streams
|
||||
t.streams = map[uint32]*agentStream{}
|
||||
t.mu.Unlock()
|
||||
for _, s := range streams {
|
||||
s.close()
|
||||
}
|
||||
t.conn.Close()
|
||||
}
|
||||
|
||||
// --- direct UDP forwarding ---
|
||||
|
||||
// runUDP maintains the agent's UDP endpoint on the host's 8446 channel and
|
||||
// relays datagrams to/from client-side UDP ports for "udp" forward rules.
|
||||
func (a *Agent) runUDP(ctx context.Context) {
|
||||
backoff := time.Second
|
||||
for ctx.Err() == nil {
|
||||
if err := a.udpOnce(ctx); err != nil {
|
||||
select {
|
||||
case <-time.After(backoff):
|
||||
case <-ctx.Done():
|
||||
return
|
||||
}
|
||||
if backoff < 15*time.Second {
|
||||
backoff *= 2
|
||||
}
|
||||
continue
|
||||
}
|
||||
backoff = time.Second
|
||||
}
|
||||
}
|
||||
|
||||
func (a *Agent) udpOnce(ctx context.Context) error {
|
||||
raddr, err := net.ResolveUDPAddr("udp", a.cfg.TunnelAddr)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
conn, err := net.DialUDP("udp", nil, raddr)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer conn.Close()
|
||||
|
||||
key := tunnel.AuthKey(a.state.Token)
|
||||
register := func() {
|
||||
m := time.Now().Unix() / 60
|
||||
sig := tunnel.Sign(key, []byte(a.state.AgentID+"|"+itoa64(m)))
|
||||
_, _ = conn.Write([]byte("R" + a.state.AgentID + " " + sig))
|
||||
}
|
||||
register()
|
||||
|
||||
// Keepalive registration so the host keeps our NAT mapping fresh.
|
||||
regCtx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
go func() {
|
||||
t := time.NewTicker(20 * time.Second)
|
||||
defer t.Stop()
|
||||
for {
|
||||
select {
|
||||
case <-regCtx.Done():
|
||||
return
|
||||
case <-t.C:
|
||||
register()
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
uc := &udpClient{server: conn, dials: map[uint32]*net.UDPConn{}}
|
||||
defer uc.shutdown()
|
||||
|
||||
buf := make([]byte, 65535)
|
||||
for {
|
||||
n, err := conn.Read(buf)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if n < 7 || buf[0] != 'D' {
|
||||
continue
|
||||
}
|
||||
sid := binary.BigEndian.Uint32(buf[1:5])
|
||||
port := int(binary.BigEndian.Uint16(buf[5:7]))
|
||||
uc.send(sid, port, append([]byte(nil), buf[7:n]...))
|
||||
}
|
||||
}
|
||||
|
||||
type udpClient struct {
|
||||
server *net.UDPConn // connection back to host:8446
|
||||
mu sync.Mutex
|
||||
dials map[uint32]*net.UDPConn // sid -> client-side socket
|
||||
}
|
||||
|
||||
// send forwards a datagram to 127.0.0.1:port and ensures replies are relayed.
|
||||
func (u *udpClient) send(sid uint32, port int, data []byte) {
|
||||
u.mu.Lock()
|
||||
c := u.dials[sid]
|
||||
u.mu.Unlock()
|
||||
if c == nil {
|
||||
raddr, err := net.ResolveUDPAddr("udp", net.JoinHostPort("127.0.0.1", itoa(port)))
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
c, err = net.DialUDP("udp", nil, raddr)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
u.mu.Lock()
|
||||
u.dials[sid] = c
|
||||
u.mu.Unlock()
|
||||
go u.replyLoop(sid, c)
|
||||
}
|
||||
_, _ = c.Write(data)
|
||||
}
|
||||
|
||||
func (u *udpClient) replyLoop(sid uint32, c *net.UDPConn) {
|
||||
buf := make([]byte, 65535)
|
||||
hdr := make([]byte, 7)
|
||||
hdr[0] = 'D'
|
||||
binary.BigEndian.PutUint32(hdr[1:5], sid)
|
||||
for {
|
||||
_ = c.SetReadDeadline(time.Now().Add(90 * time.Second))
|
||||
n, err := c.Read(buf)
|
||||
if err != nil {
|
||||
u.mu.Lock()
|
||||
delete(u.dials, sid)
|
||||
u.mu.Unlock()
|
||||
c.Close()
|
||||
return
|
||||
}
|
||||
out := append(append([]byte(nil), hdr...), buf[:n]...)
|
||||
_, _ = u.server.Write(out)
|
||||
}
|
||||
}
|
||||
|
||||
func (u *udpClient) shutdown() {
|
||||
u.mu.Lock()
|
||||
dials := u.dials
|
||||
u.dials = map[uint32]*net.UDPConn{}
|
||||
u.mu.Unlock()
|
||||
for _, c := range dials {
|
||||
c.Close()
|
||||
}
|
||||
}
|
||||
|
||||
// --- small helpers ---
|
||||
|
||||
func readLineConn(c net.Conn, max int) (string, error) {
|
||||
var b []byte
|
||||
one := make([]byte, 1)
|
||||
for len(b) < max {
|
||||
n, err := c.Read(one)
|
||||
if n > 0 {
|
||||
if one[0] == '\n' {
|
||||
return string(b), nil
|
||||
}
|
||||
b = append(b, one[0])
|
||||
}
|
||||
if err != nil {
|
||||
return "", err
|
||||
}
|
||||
}
|
||||
return string(b), nil
|
||||
}
|
||||
|
||||
func hexDecode(s string) ([]byte, error) {
|
||||
out := make([]byte, len(s)/2)
|
||||
for i := 0; i < len(out); i++ {
|
||||
hi, err := hexVal(s[i*2])
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
lo, err := hexVal(s[i*2+1])
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
out[i] = hi<<4 | lo
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func hexVal(c byte) (byte, error) {
|
||||
switch {
|
||||
case c >= '0' && c <= '9':
|
||||
return c - '0', nil
|
||||
case c >= 'a' && c <= 'f':
|
||||
return c - 'a' + 10, nil
|
||||
case c >= 'A' && c <= 'F':
|
||||
return c - 'A' + 10, nil
|
||||
}
|
||||
return 0, fmt.Errorf("bad hex")
|
||||
}
|
||||
Reference in new issue
Block a user