506 lines
13 KiB
Go
506 lines
13 KiB
Go
package main
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"fmt"
|
|
"net"
|
|
"strconv"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"dragontcp/internal/protocol"
|
|
)
|
|
|
|
type chunkSession struct {
|
|
id string
|
|
target net.Conn
|
|
maxChunk int
|
|
maxChunks int
|
|
|
|
mu sync.Mutex
|
|
notify chan struct{}
|
|
chunks map[uint64][]byte
|
|
nextDown uint64
|
|
eof bool
|
|
closed bool
|
|
lastSeen time.Time
|
|
debug *serverDebug
|
|
|
|
upMu sync.Mutex
|
|
expectedUp uint64
|
|
lastUpSeq uint64
|
|
lastUpLen int
|
|
haveLastUp bool
|
|
}
|
|
|
|
func newChunkSession(id string, target net.Conn, maxChunk, maxChunks int, debug *serverDebug) *chunkSession {
|
|
s := &chunkSession{
|
|
id: id,
|
|
target: target,
|
|
maxChunk: maxChunk,
|
|
maxChunks: maxChunks,
|
|
notify: make(chan struct{}),
|
|
chunks: make(map[uint64][]byte, maxChunks),
|
|
lastSeen: time.Now(),
|
|
debug: debug,
|
|
}
|
|
go s.readTarget()
|
|
return s
|
|
}
|
|
|
|
func (s *chunkSession) signalLocked() {
|
|
close(s.notify)
|
|
s.notify = make(chan struct{})
|
|
}
|
|
|
|
func (s *chunkSession) touchLocked() {
|
|
s.lastSeen = time.Now()
|
|
}
|
|
|
|
func (s *chunkSession) touch() {
|
|
s.mu.Lock()
|
|
s.touchLocked()
|
|
s.mu.Unlock()
|
|
}
|
|
|
|
func (s *chunkSession) readTarget() {
|
|
buf := make([]byte, s.maxChunk)
|
|
|
|
for {
|
|
n, err := s.target.Read(buf)
|
|
if n > 0 {
|
|
data := append([]byte(nil), buf[:n]...)
|
|
if s.debug != nil && s.debug.enabled {
|
|
s.debug.bytesDown.Add(uint64(n))
|
|
}
|
|
|
|
for {
|
|
s.mu.Lock()
|
|
if s.closed {
|
|
s.mu.Unlock()
|
|
return
|
|
}
|
|
if len(s.chunks) < s.maxChunks {
|
|
seq := s.nextDown
|
|
s.nextDown++
|
|
s.chunks[seq] = data
|
|
s.touchLocked()
|
|
s.signalLocked()
|
|
s.mu.Unlock()
|
|
break
|
|
}
|
|
ch := s.notify
|
|
s.mu.Unlock()
|
|
<-ch
|
|
}
|
|
}
|
|
|
|
if err != nil {
|
|
if s.debug != nil && s.debug.enabled {
|
|
s.debug.logf("TARGET EOF session=%s err=%v", s.id, err)
|
|
}
|
|
s.mu.Lock()
|
|
if !s.closed {
|
|
s.eof = true
|
|
s.touchLocked()
|
|
s.signalLocked()
|
|
}
|
|
s.mu.Unlock()
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
// push is idempotent for the most recently accepted sequence. This matters
|
|
// when the server receives a record but the tiny ACK is lost: the client can
|
|
// retry the same sequence at a smaller adaptive size without duplicating bytes
|
|
// in the target stream. The ACK reports the length that was actually accepted.
|
|
func (s *chunkSession) push(seq uint64, data []byte) (int, error) {
|
|
s.upMu.Lock()
|
|
defer s.upMu.Unlock()
|
|
|
|
if len(data) == 0 || len(data) > s.maxChunk {
|
|
return 0, fmt.Errorf("upload record size %d is invalid", len(data))
|
|
}
|
|
|
|
if s.haveLastUp && seq == s.lastUpSeq {
|
|
s.touch()
|
|
return s.lastUpLen, nil
|
|
}
|
|
|
|
if seq < s.expectedUp {
|
|
return 0, fmt.Errorf("upload sequence %d is too old", seq)
|
|
}
|
|
if seq > s.expectedUp {
|
|
return 0, fmt.Errorf("unexpected upload sequence %d, expected %d", seq, s.expectedUp)
|
|
}
|
|
|
|
if _, err := s.target.Write(data); err != nil {
|
|
return 0, err
|
|
}
|
|
|
|
if s.debug != nil && s.debug.enabled {
|
|
s.debug.bytesUp.Add(uint64(len(data)))
|
|
s.debug.pushRecords.Add(1)
|
|
}
|
|
|
|
s.lastUpSeq = seq
|
|
s.lastUpLen = len(data)
|
|
s.haveLastUp = true
|
|
s.expectedUp++
|
|
s.touch()
|
|
return len(data), nil
|
|
}
|
|
|
|
// pull returns at most limit bytes from the requested stored chunk, beginning
|
|
// at offset. The chunk sequence stays stable while the client retries smaller
|
|
// fragments, so a large queued chunk can always be recovered after an MTU-like
|
|
// failure without reopening the proxied destination connection.
|
|
func (s *chunkSession) pull(want uint64, ack int64, offset, limit int, wait time.Duration) (data []byte, total int, eof bool, final uint64, waitExpired bool, err error) {
|
|
if offset < 0 || limit <= 0 || limit > s.maxChunk {
|
|
return nil, 0, false, 0, false, fmt.Errorf("invalid pull offset/limit")
|
|
}
|
|
|
|
timer := time.NewTimer(wait)
|
|
defer timer.Stop()
|
|
|
|
for {
|
|
s.mu.Lock()
|
|
s.touchLocked()
|
|
|
|
if ack >= 0 {
|
|
removed := false
|
|
for seq := range s.chunks {
|
|
if seq <= uint64(ack) {
|
|
delete(s.chunks, seq)
|
|
removed = true
|
|
}
|
|
}
|
|
if removed {
|
|
s.signalLocked()
|
|
}
|
|
}
|
|
|
|
if chunk, ok := s.chunks[want]; ok {
|
|
if offset >= len(chunk) {
|
|
s.mu.Unlock()
|
|
return nil, len(chunk), false, 0, false, fmt.Errorf("pull offset %d beyond chunk size %d", offset, len(chunk))
|
|
}
|
|
end := offset + limit
|
|
if end > len(chunk) {
|
|
end = len(chunk)
|
|
}
|
|
out := append([]byte(nil), chunk[offset:end]...)
|
|
total = len(chunk)
|
|
s.mu.Unlock()
|
|
return out, total, false, 0, false, nil
|
|
}
|
|
|
|
if s.eof && want >= s.nextDown {
|
|
final = s.nextDown
|
|
s.mu.Unlock()
|
|
return nil, 0, true, final, false, nil
|
|
}
|
|
|
|
if s.closed {
|
|
final = s.nextDown
|
|
s.mu.Unlock()
|
|
return nil, 0, true, final, false, nil
|
|
}
|
|
|
|
ch := s.notify
|
|
s.mu.Unlock()
|
|
|
|
select {
|
|
case <-ch:
|
|
continue
|
|
case <-timer.C:
|
|
return nil, 0, false, 0, true, nil
|
|
}
|
|
}
|
|
}
|
|
|
|
func (s *chunkSession) close() {
|
|
s.mu.Lock()
|
|
if s.closed {
|
|
s.mu.Unlock()
|
|
return
|
|
}
|
|
s.closed = true
|
|
s.signalLocked()
|
|
s.mu.Unlock()
|
|
_ = s.target.Close()
|
|
}
|
|
|
|
type chunkManager struct {
|
|
mu sync.RWMutex
|
|
sessions map[string]*chunkSession
|
|
timeout time.Duration
|
|
debug *serverDebug
|
|
}
|
|
|
|
func newChunkManager(timeout time.Duration, debug *serverDebug) *chunkManager {
|
|
m := &chunkManager{
|
|
sessions: make(map[string]*chunkSession),
|
|
timeout: timeout,
|
|
debug: debug,
|
|
}
|
|
go m.cleanupLoop()
|
|
return m
|
|
}
|
|
|
|
func (m *chunkManager) get(id string) *chunkSession {
|
|
m.mu.RLock()
|
|
s := m.sessions[id]
|
|
m.mu.RUnlock()
|
|
return s
|
|
}
|
|
|
|
func (m *chunkManager) count() int {
|
|
m.mu.RLock()
|
|
n := len(m.sessions)
|
|
m.mu.RUnlock()
|
|
return n
|
|
}
|
|
|
|
func (m *chunkManager) add(id string, s *chunkSession) error {
|
|
m.mu.Lock()
|
|
defer m.mu.Unlock()
|
|
if _, exists := m.sessions[id]; exists {
|
|
return fmt.Errorf("session already exists")
|
|
}
|
|
m.sessions[id] = s
|
|
return nil
|
|
}
|
|
|
|
func (m *chunkManager) remove(id string) {
|
|
m.mu.Lock()
|
|
s := m.sessions[id]
|
|
delete(m.sessions, id)
|
|
m.mu.Unlock()
|
|
if s != nil {
|
|
s.close()
|
|
}
|
|
}
|
|
|
|
func (m *chunkManager) cleanupLoop() {
|
|
ticker := time.NewTicker(30 * time.Second)
|
|
defer ticker.Stop()
|
|
|
|
for range ticker.C {
|
|
cutoff := time.Now().Add(-m.timeout)
|
|
var stale []string
|
|
|
|
m.mu.RLock()
|
|
for id, s := range m.sessions {
|
|
s.mu.Lock()
|
|
last := s.lastSeen
|
|
closed := s.closed
|
|
s.mu.Unlock()
|
|
if closed || last.Before(cutoff) {
|
|
stale = append(stale, id)
|
|
}
|
|
}
|
|
m.mu.RUnlock()
|
|
|
|
for _, id := range stale {
|
|
if m.debug != nil && m.debug.enabled {
|
|
m.debug.logf("SESSION timeout-close id=%s active_sessions=%d", id, m.count())
|
|
}
|
|
m.remove(id)
|
|
if m.debug != nil && m.debug.enabled {
|
|
m.debug.sessionsClosed.Add(1)
|
|
m.debug.activeSessions.Add(-1)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func decodeWireToken(token string) string {
|
|
if token == "-" {
|
|
return ""
|
|
}
|
|
return token
|
|
}
|
|
|
|
func isChunkCommand(payload []byte) bool {
|
|
return bytes.HasPrefix(payload, []byte("COPEN ")) ||
|
|
bytes.HasPrefix(payload, []byte("CPUSH ")) ||
|
|
bytes.HasPrefix(payload, []byte("CPULL ")) ||
|
|
bytes.HasPrefix(payload, []byte("CCLOSE "))
|
|
}
|
|
|
|
func processChunkCommand(
|
|
conn net.Conn,
|
|
requestID uint32,
|
|
payload []byte,
|
|
token string,
|
|
allowPrivate bool,
|
|
cache *dnsCache,
|
|
tcpBuffer int,
|
|
manager *chunkManager,
|
|
maxChunk int,
|
|
maxBufferedChunks int,
|
|
pollWait time.Duration,
|
|
debug *serverDebug,
|
|
) error {
|
|
if bytes.HasPrefix(payload, []byte("COPEN ")) {
|
|
parts := strings.Fields(string(payload))
|
|
if len(parts) != 5 {
|
|
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR bad COPEN"))
|
|
}
|
|
if !tokenEqual(decodeWireToken(parts[1]), token) {
|
|
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR authentication failed"))
|
|
}
|
|
sid := parts[2]
|
|
if len(sid) < 16 || len(sid) > 64 {
|
|
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR invalid session id"))
|
|
}
|
|
host := parts[3]
|
|
port, err := strconv.Atoi(parts[4])
|
|
if err != nil || port < 1 || port > 65535 {
|
|
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR invalid port"))
|
|
}
|
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
|
target, err := dialTarget(ctx, host, port, allowPrivate, cache, tcpBuffer)
|
|
cancel()
|
|
if err != nil {
|
|
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR "+err.Error()))
|
|
}
|
|
|
|
session := newChunkSession(sid, target, maxChunk, maxBufferedChunks, debug)
|
|
if err := manager.add(sid, session); err != nil {
|
|
session.close()
|
|
if debug != nil && debug.enabled {
|
|
debug.errorf("COPEN session=%s target=%s:%d failed: %v", sid, host, port, err)
|
|
}
|
|
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR "+err.Error()))
|
|
}
|
|
if debug != nil && debug.enabled {
|
|
debug.sessionsOpened.Add(1)
|
|
debug.activeSessions.Add(1)
|
|
debug.logf("SESSION OPEN id=%s peer=%v target=%s:%d max_chunk=%d active_sessions=%d", sid, conn.RemoteAddr(), host, port, maxChunk, manager.count())
|
|
debug.chunkf("COPEN id=%s target=%s:%d -> OPENED max=%d", sid, host, port, maxChunk)
|
|
}
|
|
return protocol.WriteResponseFrame(conn, requestID, []byte(fmt.Sprintf("OPENED %d", maxChunk)))
|
|
}
|
|
|
|
if bytes.HasPrefix(payload, []byte("CPUSH ")) {
|
|
parts := bytes.SplitN(payload, []byte(" "), 5)
|
|
if len(parts) != 5 {
|
|
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR bad CPUSH"))
|
|
}
|
|
if !tokenEqual(decodeWireToken(string(parts[1])), token) {
|
|
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR authentication failed"))
|
|
}
|
|
sid := string(parts[2])
|
|
seq, err := strconv.ParseUint(string(parts[3]), 10, 64)
|
|
if err != nil {
|
|
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR invalid sequence"))
|
|
}
|
|
s := manager.get(sid)
|
|
if s == nil {
|
|
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR unknown session"))
|
|
}
|
|
accepted, err := s.push(seq, parts[4])
|
|
if err != nil {
|
|
if debug != nil && debug.enabled {
|
|
debug.errorf("CPUSH id=%s seq=%d bytes=%d: %v", sid, seq, len(parts[4]), err)
|
|
}
|
|
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR "+err.Error()))
|
|
}
|
|
if debug != nil {
|
|
debug.chunkf("CPUSH id=%s seq=%d bytes=%d -> ACK accepted=%d", sid, seq, len(parts[4]), accepted)
|
|
}
|
|
return protocol.WriteResponseFrame(conn, requestID, []byte(fmt.Sprintf("ACK %d %d", seq, accepted)))
|
|
}
|
|
|
|
if bytes.HasPrefix(payload, []byte("CPULL ")) {
|
|
parts := strings.Fields(string(payload))
|
|
if len(parts) != 7 {
|
|
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR bad CPULL"))
|
|
}
|
|
if !tokenEqual(decodeWireToken(parts[1]), token) {
|
|
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR authentication failed"))
|
|
}
|
|
s := manager.get(parts[2])
|
|
if s == nil {
|
|
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR unknown session"))
|
|
}
|
|
ack, err := strconv.ParseInt(parts[3], 10, 64)
|
|
if err != nil || ack < -1 {
|
|
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR invalid ack"))
|
|
}
|
|
want, err := strconv.ParseUint(parts[4], 10, 64)
|
|
if err != nil {
|
|
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR invalid want"))
|
|
}
|
|
offset, err := strconv.Atoi(parts[5])
|
|
if err != nil || offset < 0 {
|
|
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR invalid offset"))
|
|
}
|
|
limit, err := strconv.Atoi(parts[6])
|
|
if err != nil || limit < 1 {
|
|
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR invalid limit"))
|
|
}
|
|
if limit > maxChunk {
|
|
limit = maxChunk
|
|
}
|
|
if debug != nil && debug.enabled {
|
|
debug.pullRequests.Add(1)
|
|
debug.chunkf("CPULL id=%s ack=%d want=%d offset=%d limit=%d", parts[2], ack, want, offset, limit)
|
|
}
|
|
|
|
data, total, eof, final, waitExpired, err := s.pull(want, ack, offset, limit, pollWait)
|
|
if err != nil {
|
|
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR "+err.Error()))
|
|
}
|
|
if waitExpired {
|
|
if debug != nil && debug.enabled {
|
|
debug.waitRecords.Add(1)
|
|
debug.chunkf("CPULL id=%s want=%d -> WAIT", parts[2], want)
|
|
}
|
|
return protocol.WriteResponseFrame(conn, requestID, []byte("WAIT"))
|
|
}
|
|
if eof {
|
|
if debug != nil {
|
|
debug.chunkf("CPULL id=%s want=%d -> EOF final=%d", parts[2], want, final)
|
|
}
|
|
return protocol.WriteResponseFrame(conn, requestID, []byte(fmt.Sprintf("EOF %d", final)))
|
|
}
|
|
|
|
if debug != nil && debug.enabled {
|
|
debug.dataRecords.Add(1)
|
|
debug.chunkf("DATA id=%s seq=%d offset=%d bytes=%d total=%d", parts[2], want, offset, len(data), total)
|
|
}
|
|
prefix := []byte(fmt.Sprintf("DATA %d %d %d ", want, offset, total))
|
|
out := make([]byte, len(prefix)+len(data))
|
|
copy(out, prefix)
|
|
copy(out[len(prefix):], data)
|
|
return protocol.WriteResponseFrame(conn, requestID, out)
|
|
}
|
|
|
|
if bytes.HasPrefix(payload, []byte("CCLOSE ")) {
|
|
parts := strings.Fields(string(payload))
|
|
if len(parts) != 3 {
|
|
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR bad CCLOSE"))
|
|
}
|
|
if !tokenEqual(decodeWireToken(parts[1]), token) {
|
|
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR authentication failed"))
|
|
}
|
|
manager.remove(parts[2])
|
|
if debug != nil && debug.enabled {
|
|
debug.sessionsClosed.Add(1)
|
|
debug.activeSessions.Add(-1)
|
|
debug.logf("SESSION CLOSE id=%s peer=%v active_sessions=%d", parts[2], conn.RemoteAddr(), manager.count())
|
|
debug.chunkf("CCLOSE id=%s -> CLOSED", parts[2])
|
|
}
|
|
return protocol.WriteResponseFrame(conn, requestID, []byte("CLOSED"))
|
|
}
|
|
|
|
return protocol.WriteResponseFrame(conn, requestID, []byte("ERR unknown chunk command"))
|
|
}
|