syncthing/session.go
2015-08-20 14:29:57 +02:00

259 lines
5.2 KiB
Go

// Copyright (C) 2015 Audrius Butkevicius and Contributors (see the CONTRIBUTORS file).
package main
import (
"crypto/rand"
"encoding/hex"
"fmt"
"log"
"net"
"sync"
"sync/atomic"
"time"
"github.com/juju/ratelimit"
"github.com/syncthing/relaysrv/protocol"
syncthingprotocol "github.com/syncthing/protocol"
)
var (
sessionMut = sync.Mutex{}
sessions = make(map[string]*session, 0)
numProxies int64
bytesProxied int64
)
type session struct {
serverkey []byte
clientkey []byte
rateLimit func(bytes int64)
conns chan net.Conn
}
func newSession(sessionRateLimit, globalRateLimit *ratelimit.Bucket) *session {
serverkey := make([]byte, 32)
_, err := rand.Read(serverkey)
if err != nil {
return nil
}
clientkey := make([]byte, 32)
_, err = rand.Read(clientkey)
if err != nil {
return nil
}
ses := &session{
serverkey: serverkey,
clientkey: clientkey,
rateLimit: makeRateLimitFunc(sessionRateLimit, globalRateLimit),
conns: make(chan net.Conn),
}
if debug {
log.Println("New session", ses)
}
sessionMut.Lock()
sessions[string(ses.serverkey)] = ses
sessions[string(ses.clientkey)] = ses
sessionMut.Unlock()
return ses
}
func findSession(key string) *session {
sessionMut.Lock()
defer sessionMut.Unlock()
lob, ok := sessions[key]
if !ok {
return nil
}
delete(sessions, key)
return lob
}
func (s *session) AddConnection(conn net.Conn) bool {
if debug {
log.Println("New connection for", s, "from", conn.RemoteAddr())
}
select {
case s.conns <- conn:
return true
default:
}
return false
}
func (s *session) Serve() {
timedout := time.After(messageTimeout)
if debug {
log.Println("Session", s, "serving")
}
conns := make([]net.Conn, 0, 2)
for {
select {
case conn := <-s.conns:
conns = append(conns, conn)
if len(conns) < 2 {
continue
}
close(s.conns)
if debug {
log.Println("Session", s, "starting between", conns[0].RemoteAddr(), "and", conns[1].RemoteAddr())
}
wg := sync.WaitGroup{}
wg.Add(2)
var err0 error
go func() {
err0 = s.proxy(conns[0], conns[1])
wg.Done()
}()
var err1 error
go func() {
err1 = s.proxy(conns[1], conns[0])
wg.Done()
}()
wg.Wait()
if debug {
log.Println("Session", s, "ended, outcomes:", err0, "and", err1)
}
goto done
case <-timedout:
if debug {
log.Println("Session", s, "timed out")
}
goto done
}
}
done:
sessionMut.Lock()
delete(sessions, string(s.serverkey))
delete(sessions, string(s.clientkey))
sessionMut.Unlock()
for _, conn := range conns {
conn.Close()
}
if debug {
log.Println("Session", s, "stopping")
}
}
func (s *session) GetClientInvitationMessage(from syncthingprotocol.DeviceID) protocol.SessionInvitation {
return protocol.SessionInvitation{
From: from[:],
Key: []byte(s.clientkey),
Address: sessionAddress,
Port: sessionPort,
ServerSocket: false,
}
}
func (s *session) GetServerInvitationMessage(from syncthingprotocol.DeviceID) protocol.SessionInvitation {
return protocol.SessionInvitation{
From: from[:],
Key: []byte(s.serverkey),
Address: sessionAddress,
Port: sessionPort,
ServerSocket: true,
}
}
func (s *session) proxy(c1, c2 net.Conn) error {
if debug {
log.Println("Proxy", c1.RemoteAddr(), "->", c2.RemoteAddr())
}
atomic.AddInt64(&numProxies, 1)
defer atomic.AddInt64(&numProxies, -1)
buf := make([]byte, 65536)
for {
c1.SetReadDeadline(time.Now().Add(networkTimeout))
n, err := c1.Read(buf)
if err != nil {
return err
}
atomic.AddInt64(&bytesProxied, int64(n))
if debug {
log.Printf("%d bytes from %s to %s", n, c1.RemoteAddr(), c2.RemoteAddr())
}
if s.rateLimit != nil {
s.rateLimit(int64(n))
}
c2.SetWriteDeadline(time.Now().Add(networkTimeout))
_, err = c2.Write(buf[:n])
if err != nil {
return err
}
}
}
func (s *session) String() string {
return fmt.Sprintf("<%s/%s>", hex.EncodeToString(s.clientkey)[:5], hex.EncodeToString(s.serverkey)[:5])
}
func makeRateLimitFunc(sessionRateLimit, globalRateLimit *ratelimit.Bucket) func(int64) {
// This may be a case of super duper premature optimization... We build an
// optimized function to do the rate limiting here based on what we need
// to do and then use it in the loop.
if sessionRateLimit == nil && globalRateLimit == nil {
// No limiting needed. We could equally well return a func(int64){} and
// not do a nil check were we use it, but I think the nil check there
// makes it clear that there will be no limiting if none is
// configured...
return nil
}
if sessionRateLimit == nil {
// We only have a global limiter
return func(bytes int64) {
globalRateLimit.Wait(bytes)
}
}
if globalRateLimit == nil {
// We only have a session limiter
return func(bytes int64) {
sessionRateLimit.Wait(bytes)
}
}
// We have both. Queue the bytes on both the global and session specific
// rate limiters. Wait for both in parallell, so that the actual send
// happens when both conditions are satisfied. In practice this just means
// wait the longer of the two times.
return func(bytes int64) {
t0 := sessionRateLimit.Take(bytes)
t1 := globalRateLimit.Take(bytes)
if t0 > t1 {
time.Sleep(t0)
} else {
time.Sleep(t1)
}
}
}