Cloak/internal/multiplex/switchboard.go

173 lines
4.1 KiB
Go
Raw Normal View History

2018-10-05 22:44:20 +00:00
package multiplex
import (
2018-10-28 22:51:00 +00:00
"errors"
2019-08-02 14:45:33 +00:00
log "github.com/sirupsen/logrus"
"math/rand"
2018-10-05 22:44:20 +00:00
"net"
2018-10-28 21:22:38 +00:00
"sync"
"sync/atomic"
2018-10-05 22:44:20 +00:00
)
2019-08-12 21:43:16 +00:00
const (
FIXED_CONN_MAPPING switchboardStrategy = iota
UNIFORM_SPREAD
)
type switchboardConfig struct {
2020-04-08 14:15:08 +00:00
valve Valve
strategy switchboardStrategy
recvBufferSize int
}
2019-09-01 22:14:28 +00:00
// switchboard is responsible for keeping the reference of TCP connections between client and server
2018-10-05 22:44:20 +00:00
type switchboard struct {
session *Session
switchboardConfig
2018-10-28 22:51:00 +00:00
conns sync.Map
2020-04-17 00:04:26 +00:00
numConns uint32
nextConnId uint32
2019-08-05 13:33:20 +00:00
broken uint32
2018-11-07 21:16:13 +00:00
}
func makeSwitchboard(sesh *Session, config switchboardConfig) *switchboard {
2018-11-07 21:16:13 +00:00
// rates are uint64 because in the usermanager we want the bandwidth to be atomically
// operated (so that the bandwidth can change on the fly).
2018-10-05 22:44:20 +00:00
sb := &switchboard{
session: sesh,
switchboardConfig: config,
2020-04-12 11:51:00 +00:00
nextConnId: 1,
2018-10-05 22:44:20 +00:00
}
return sb
}
2019-08-05 13:33:20 +00:00
var errBrokenSwitchboard = errors.New("the switchboard is broken")
2018-10-28 22:51:00 +00:00
func (sb *switchboard) connsCount() int {
2020-04-17 00:04:26 +00:00
return int(atomic.LoadUint32(&sb.numConns))
}
func (sb *switchboard) addConn(conn net.Conn) {
connId := atomic.AddUint32(&sb.nextConnId, 1) - 1
2020-04-17 00:04:26 +00:00
atomic.AddUint32(&sb.numConns, 1)
sb.conns.Store(connId, conn)
go sb.deplex(connId, conn)
2018-10-28 21:22:38 +00:00
}
2018-10-20 16:03:39 +00:00
// a pointer to connId is passed here so that the switchboard can reassign it
2019-08-14 09:44:50 +00:00
func (sb *switchboard) send(data []byte, connId *uint32) (n int, err error) {
2019-10-20 14:37:41 +00:00
writeAndRegUsage := func(conn net.Conn, d []byte) (int, error) {
n, err = conn.Write(d)
if err != nil {
2020-04-17 00:04:26 +00:00
sb.conns.Delete(*connId)
2019-10-20 14:37:41 +00:00
sb.close("failed to write to remote " + err.Error())
return n, err
}
2020-04-08 14:15:08 +00:00
sb.valve.AddTx(int64(n))
2019-10-20 14:37:41 +00:00
return n, nil
}
2020-04-08 14:15:08 +00:00
sb.valve.txWait(len(data))
2020-04-09 18:10:04 +00:00
if atomic.LoadUint32(&sb.broken) == 1 || sb.connsCount() == 0 {
2019-10-20 14:37:41 +00:00
return 0, errBrokenSwitchboard
}
2019-09-22 09:44:16 +00:00
2020-04-03 21:46:07 +00:00
switch sb.strategy {
case UNIFORM_SPREAD:
_, conn, err := sb.pickRandConn()
if err != nil {
return 0, errBrokenSwitchboard
2019-09-22 09:44:16 +00:00
}
return writeAndRegUsage(conn, data)
2020-04-03 21:46:07 +00:00
case FIXED_CONN_MAPPING:
connI, ok := sb.conns.Load(*connId)
2019-08-12 21:43:16 +00:00
if ok {
conn := connI.(net.Conn)
2019-10-20 14:37:41 +00:00
return writeAndRegUsage(conn, data)
2019-08-12 21:43:16 +00:00
} else {
newConnId, conn, err := sb.pickRandConn()
if err != nil {
return 0, errBrokenSwitchboard
2019-08-21 12:14:23 +00:00
}
*connId = newConnId
return writeAndRegUsage(conn, data)
}
2020-04-03 21:46:07 +00:00
default:
return 0, errors.New("unsupported traffic distribution strategy")
2018-10-28 21:22:38 +00:00
}
2018-10-05 22:44:20 +00:00
}
2019-08-20 21:43:04 +00:00
// returns a random connId
func (sb *switchboard) pickRandConn() (uint32, net.Conn, error) {
connCount := sb.connsCount()
if atomic.LoadUint32(&sb.broken) == 1 || connCount == 0 {
return 0, nil, errBrokenSwitchboard
2018-10-28 21:22:38 +00:00
}
2019-08-31 14:32:57 +00:00
// there is no guarantee that sb.conns still has the same amount of entries
// between the count loop and the pick loop
// so if the r > len(sb.conns) at the point of range call, the last visited element is picked
var id uint32
var conn net.Conn
r := rand.Intn(connCount)
2020-01-22 22:17:18 +00:00
var c int
sb.conns.Range(func(connIdI, connI interface{}) bool {
2019-08-31 14:32:57 +00:00
if r == c {
id = connIdI.(uint32)
conn = connI.(net.Conn)
return false
2019-08-31 14:32:57 +00:00
}
c++
return true
})
2020-04-03 21:46:07 +00:00
// if len(sb.conns) is 0
if conn == nil {
return 0, nil, errBrokenSwitchboard
}
return id, conn, nil
2018-10-28 21:22:38 +00:00
}
2019-10-20 14:37:41 +00:00
func (sb *switchboard) close(terminalMsg string) {
atomic.StoreUint32(&sb.broken, 1)
if !sb.session.IsClosed() {
sb.session.SetTerminalMsg(terminalMsg)
sb.session.passiveClose()
}
}
2018-12-30 01:15:46 +00:00
// actively triggered by session.Close()
2018-12-31 11:30:39 +00:00
func (sb *switchboard) closeAll() {
sb.conns.Range(func(key, connI interface{}) bool {
conn := connI.(net.Conn)
conn.Close()
sb.conns.Delete(key)
return true
})
2018-10-05 22:44:20 +00:00
}
2019-08-05 13:33:20 +00:00
// deplex function costantly reads from a TCP connection
func (sb *switchboard) deplex(connId uint32, conn net.Conn) {
2020-04-05 00:15:45 +00:00
defer conn.Close()
buf := make([]byte, sb.recvBufferSize)
2018-10-09 20:53:55 +00:00
for {
n, err := conn.Read(buf)
2020-04-08 14:15:08 +00:00
sb.valve.rxWait(n)
sb.valve.AddRx(int64(n))
2018-10-09 20:53:55 +00:00
if err != nil {
2019-08-31 14:32:57 +00:00
log.Debugf("a connection for session %v has closed: %v", sb.session.id, err)
2020-04-12 11:51:00 +00:00
sb.conns.Delete(connId)
2020-04-17 00:04:26 +00:00
atomic.AddUint32(&sb.numConns, ^uint32(0))
2020-04-08 15:41:39 +00:00
sb.close("a connection has dropped unexpectedly")
2018-10-09 20:53:55 +00:00
return
}
2019-08-16 22:47:15 +00:00
err = sb.session.recvDataFromRemote(buf[:n])
if err != nil {
log.Error(err)
}
2018-10-05 22:44:20 +00:00
}
}