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,246 @@
|
||||
package server
|
||||
|
||||
import (
|
||||
"encoding/binary"
|
||||
"log"
|
||||
"net"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/x-grk/longpolling-remote-access/internal/tunnel"
|
||||
)
|
||||
|
||||
// udpRelay carries UDP forwards over the direct 8446 UDP channel. The agent
|
||||
// registers its UDP return address; end-user datagrams arriving on a rule's
|
||||
// host port are relayed to the agent and replies are sent back.
|
||||
//
|
||||
// Datagram framing on the 8446 control socket:
|
||||
//
|
||||
// registration (agent->server): "R <agentID> <sig>"
|
||||
// data (both): 'D' | sid(4) | port(2) | payload
|
||||
type udpRelay struct {
|
||||
db *DB
|
||||
mu sync.Mutex
|
||||
ctrl *net.UDPConn
|
||||
|
||||
agentAddr map[string]*net.UDPAddr // agentID -> last registered addr
|
||||
sessions map[uint32]*udpSession
|
||||
nextSID uint32
|
||||
rulePorts map[int64]*net.UDPConn // ruleID -> per-rule listener
|
||||
}
|
||||
|
||||
type udpSession struct {
|
||||
id uint32
|
||||
agentID string
|
||||
ruleConn *net.UDPConn
|
||||
userAddr *net.UDPAddr
|
||||
clientPort int
|
||||
last time.Time
|
||||
}
|
||||
|
||||
func newUDPRelay(db *DB) *udpRelay {
|
||||
return &udpRelay{
|
||||
db: db,
|
||||
agentAddr: map[string]*net.UDPAddr{},
|
||||
sessions: map[uint32]*udpSession{},
|
||||
rulePorts: map[int64]*net.UDPConn{},
|
||||
}
|
||||
}
|
||||
|
||||
// listen opens the shared 8446 UDP control socket.
|
||||
func (u *udpRelay) listen(addr string) error {
|
||||
ua, err := net.ResolveUDPAddr("udp", addr)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
c, err := net.ListenUDP("udp", ua)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
u.ctrl = c
|
||||
log.Printf("tunnel: listening on %s/udp (direct UDP forwarding)", addr)
|
||||
go u.ctrlLoop()
|
||||
go u.gc()
|
||||
return nil
|
||||
}
|
||||
|
||||
func (u *udpRelay) ctrlLoop() {
|
||||
buf := make([]byte, 65535)
|
||||
for {
|
||||
n, src, err := u.ctrl.ReadFromUDP(buf)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if n == 0 {
|
||||
continue
|
||||
}
|
||||
switch buf[0] {
|
||||
case 'R':
|
||||
u.handleReg(string(buf[1:n]), src)
|
||||
case 'D':
|
||||
u.handleAgentData(buf[1:n])
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// handleReg validates "agentID sig" and records the agent's UDP addr.
|
||||
// sig = HMAC(key, "agentID|minute") accepting the current or previous minute.
|
||||
func (u *udpRelay) handleReg(line string, src *net.UDPAddr) {
|
||||
f := strings.Fields(strings.TrimSpace(line))
|
||||
if len(f) != 2 {
|
||||
return
|
||||
}
|
||||
agentID, sig := f[0], f[1]
|
||||
agent, err := u.db.GetAgent(agentID)
|
||||
if err != nil || agent == nil {
|
||||
return
|
||||
}
|
||||
key, err := tunnel.KeyFromHash(agent.TokenHash)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
now := time.Now().Unix() / 60
|
||||
ok := false
|
||||
for _, m := range []int64{now, now - 1} {
|
||||
msg := []byte(agentID + "|" + itoa64(m))
|
||||
if tunnel.EqualSig(sig, tunnel.Sign(key, msg)) {
|
||||
ok = true
|
||||
break
|
||||
}
|
||||
}
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
u.mu.Lock()
|
||||
u.agentAddr[agentID] = src
|
||||
u.mu.Unlock()
|
||||
}
|
||||
|
||||
// handleAgentData routes a reply datagram from an agent back to the end user.
|
||||
func (u *udpRelay) handleAgentData(b []byte) {
|
||||
if len(b) < 6 {
|
||||
return
|
||||
}
|
||||
sid := binary.BigEndian.Uint32(b[0:4])
|
||||
payload := b[6:]
|
||||
u.mu.Lock()
|
||||
s := u.sessions[sid]
|
||||
if s != nil {
|
||||
s.last = time.Now()
|
||||
}
|
||||
u.mu.Unlock()
|
||||
if s != nil {
|
||||
_, _ = s.ruleConn.WriteToUDP(payload, s.userAddr)
|
||||
}
|
||||
}
|
||||
|
||||
// openRulePort starts the per-rule UDP listener on the host port.
|
||||
func (u *udpRelay) openRulePort(rule ForwardRule) error {
|
||||
ua, err := net.ResolveUDPAddr("udp", rule.ListenAddr+":"+itoa(rule.HostPort))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
c, err := net.ListenUDP("udp", ua)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
u.mu.Lock()
|
||||
u.rulePorts[rule.ID] = c
|
||||
u.mu.Unlock()
|
||||
go u.rulePortLoop(rule, c)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (u *udpRelay) closeRulePort(ruleID int64) {
|
||||
u.mu.Lock()
|
||||
c := u.rulePorts[ruleID]
|
||||
delete(u.rulePorts, ruleID)
|
||||
u.mu.Unlock()
|
||||
if c != nil {
|
||||
c.Close()
|
||||
}
|
||||
}
|
||||
|
||||
func (u *udpRelay) rulePortLoop(rule ForwardRule, c *net.UDPConn) {
|
||||
buf := make([]byte, 65535)
|
||||
for {
|
||||
n, userAddr, err := c.ReadFromUDP(buf)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
u.mu.Lock()
|
||||
aAddr := u.agentAddr[rule.AgentID]
|
||||
// Find existing session for this (rule, user).
|
||||
var sess *udpSession
|
||||
for _, s := range u.sessions {
|
||||
if s.ruleConn == c && udpAddrEqual(s.userAddr, userAddr) {
|
||||
sess = s
|
||||
break
|
||||
}
|
||||
}
|
||||
if sess == nil {
|
||||
u.nextSID++
|
||||
sess = &udpSession{
|
||||
id: u.nextSID, agentID: rule.AgentID, ruleConn: c,
|
||||
userAddr: userAddr, clientPort: rule.ClientPort,
|
||||
}
|
||||
u.sessions[sess.id] = sess
|
||||
}
|
||||
sess.last = time.Now()
|
||||
u.mu.Unlock()
|
||||
|
||||
if aAddr == nil {
|
||||
continue // agent has not registered its UDP endpoint yet
|
||||
}
|
||||
// 'D' | sid(4) | clientPort(2) | payload
|
||||
out := make([]byte, 7+n)
|
||||
out[0] = 'D'
|
||||
binary.BigEndian.PutUint32(out[1:5], sess.id)
|
||||
binary.BigEndian.PutUint16(out[5:7], uint16(rule.ClientPort))
|
||||
copy(out[7:], buf[:n])
|
||||
_, _ = u.ctrl.WriteToUDP(out, aAddr)
|
||||
}
|
||||
}
|
||||
|
||||
func (u *udpRelay) gc() {
|
||||
t := time.NewTicker(30 * time.Second)
|
||||
defer t.Stop()
|
||||
for range t.C {
|
||||
cut := time.Now().Add(-90 * time.Second)
|
||||
u.mu.Lock()
|
||||
for id, s := range u.sessions {
|
||||
if s.last.Before(cut) {
|
||||
delete(u.sessions, id)
|
||||
}
|
||||
}
|
||||
u.mu.Unlock()
|
||||
}
|
||||
}
|
||||
|
||||
func udpAddrEqual(a, b *net.UDPAddr) bool {
|
||||
return a != nil && b != nil && a.Port == b.Port && a.IP.Equal(b.IP)
|
||||
}
|
||||
|
||||
func itoa(n int) string { return itoa64(int64(n)) }
|
||||
func itoa64(n int64) string {
|
||||
if n == 0 {
|
||||
return "0"
|
||||
}
|
||||
neg := n < 0
|
||||
if neg {
|
||||
n = -n
|
||||
}
|
||||
var b [24]byte
|
||||
i := len(b)
|
||||
for n > 0 {
|
||||
i--
|
||||
b[i] = byte('0' + n%10)
|
||||
n /= 10
|
||||
}
|
||||
if neg {
|
||||
i--
|
||||
b[i] = '-'
|
||||
}
|
||||
return string(b[i:])
|
||||
}
|
||||
Reference in new issue
Block a user