kirillius
/
Cloak
mirror of https://github.com/cbeuw/Cloak
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Redo the implementation of switchboard and remove the need for connId

This commit is contained in:
Andy Wang 2020-12-28 01:10:24 +00:00
parent 4a2eac51fe
commit 8ab0c2d96b
No known key found for this signature in database
GPG Key ID: 181B49F9F38F3374
4 changed files with 97 additions and 92 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
internal/multiplex/session.go
View File

@ -286,14 +286,12 @@ func (sesh *Session) closeSession() error {
sesh.streamsM.Lock()
close(sesh.acceptCh)
for id, stream := range sesh.streams {
if stream == nil {
continue
}
atomic.StoreUint32(&stream.closed, 1)
if stream != nil && atomic.CompareAndSwapUint32(&stream.closed, 0, 1) {
_ = stream.getRecvBuf().Close() // will not block
delete(sesh.streams, id)
sesh.streamCountDecr()
}
}
sesh.streamsM.Unlock()
return nil
}
@ -333,7 +331,7 @@ func (sesh *Session) Close() error {
if err != nil {
return err
}
_, err = sesh.sb.send((*buf)[:i], new(uint32))
_, err = sesh.sb.send((*buf)[:i], new(net.Conn))
if err != nil {
return err
}

4
internal/multiplex/stream.go
View File

@ -40,7 +40,7 @@ type Stream struct {
// recvBuffer (implemented by streamBuffer under ordered mode) will not receive out-of-order packets
// so it won't have to use its priority queue to sort it.
// This is not used in unordered connection mode
assignedConnId uint32
assignedConn net.Conn
readFromTimeout time.Duration
}
@ -119,7 +119,7 @@ func (s *Stream) obfuscateAndSend(buf []byte, payloadOffsetInBuf int) error {
return err
}
_, err = s.session.sb.send(buf[:cipherTextLen], &s.assignedConnId)
_, err = s.session.sb.send(buf[:cipherTextLen], &s.assignedConn)
if err != nil {
if err == errBrokenSwitchboard {
s.session.SetTerminalMsg(err.Error())

125
internal/multiplex/switchboard.go
View File

@ -28,10 +28,8 @@ type switchboard struct {
valve Valve
strategy switchboardStrategy
// map of connId to net.Conn
conns sync.Map
numConns uint32
nextConnId uint32
connsM sync.RWMutex
conns []net.Conn
randPool sync.Pool
broken uint32
@ -49,7 +47,6 @@ func makeSwitchboard(sesh *Session) *switchboard {
session: sesh,
strategy: strategy,
valve: sesh.Valve,
nextConnId: 1,
randPool: sync.Pool{New: func() interface{} {
return rand.New(rand.NewSource(int64(time.Now().Nanosecond())))
}},
@ -59,88 +56,95 @@ func makeSwitchboard(sesh *Session) *switchboard {
var errBrokenSwitchboard = errors.New("the switchboard is broken")
func (sb *switchboard) connsCount() int {
return int(atomic.LoadUint32(&sb.numConns))
func (sb *switchboard) delConn(conn net.Conn) {
sb.connsM.Lock()
defer sb.connsM.Unlock()
if len(sb.conns) <= 1 {
sb.conns = nil
return
}
var i int
var c net.Conn
for i, c = range sb.conns {
if c == conn {
break
}
}
sb.conns = append(sb.conns[:i], sb.conns[i+1:]...)
}
func (sb *switchboard) addConn(conn net.Conn) {
connId := atomic.AddUint32(&sb.nextConnId, 1) - 1
atomic.AddUint32(&sb.numConns, 1)
sb.conns.Store(connId, conn)
go sb.deplex(connId, conn)
sb.connsM.Lock()
sb.conns = append(sb.conns, conn)
sb.connsM.Unlock()
go sb.deplex(conn)
}
// a pointer to connId is passed here so that the switchboard can reassign it if that connId isn't usable
func (sb *switchboard) send(data []byte, connId *uint32) (n int, err error) {
// a pointer to assignedConn is passed here so that the switchboard can reassign it if that conn isn't usable
func (sb *switchboard) send(data []byte, assignedConn *net.Conn) (n int, err error) {
sb.valve.txWait(len(data))
if atomic.LoadUint32(&sb.broken) == 1 || sb.connsCount() == 0 {
if atomic.LoadUint32(&sb.broken) == 1 {
return 0, errBrokenSwitchboard
}
var conn net.Conn
switch sb.strategy {
case UNIFORM_SPREAD:
_, conn, err = sb.pickRandConn()
conn, err = sb.pickRandConn()
if err != nil {
return 0, errBrokenSwitchboard
}
case FIXED_CONN_MAPPING:
connI, ok := sb.conns.Load(*connId)
if ok {
conn = connI.(net.Conn)
} else {
var newConnId uint32
newConnId, conn, err = sb.pickRandConn()
if err != nil {
return 0, errBrokenSwitchboard
}
*connId = newConnId
}
conn = *assignedConn
default:
return 0, errors.New("unsupported traffic distribution strategy")
}
if conn != nil {
n, err = conn.Write(data)
if err != nil {
sb.conns.Delete(*connId)
sb.session.SetTerminalMsg("failed to write to remote " + err.Error())
sb.delConn(conn)
}
} else {
conn, err = sb.pickRandConn()
if err != nil {
sb.delConn(conn)
sb.session.SetTerminalMsg("failed to pick a connection " + err.Error())
sb.session.passiveClose()
return 0, err
}
n, err = conn.Write(data)
if err != nil {
sb.delConn(conn)
sb.session.SetTerminalMsg("failed to send to remote " + err.Error())
sb.session.passiveClose()
return n, err
}
*assignedConn = conn
}
sb.valve.AddTx(int64(n))
return n, nil
}
// 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
func (sb *switchboard) pickRandConn() (net.Conn, error) {
if atomic.LoadUint32(&sb.broken) == 1 {
return nil, errBrokenSwitchboard
}
// 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
randReader := sb.randPool.Get().(*rand.Rand)
r := randReader.Intn(connCount)
sb.connsM.RLock()
defer sb.connsM.RUnlock()
connsCount := len(sb.conns)
if connsCount == 0 {
return nil, errBrokenSwitchboard
}
r := randReader.Intn(connsCount)
sb.randPool.Put(randReader)
var c int
sb.conns.Range(func(connIdI, connI interface{}) bool {
if r == c {
id = connIdI.(uint32)
conn = connI.(net.Conn)
return false
}
c++
return true
})
// if len(sb.conns) is 0
if conn == nil {
return 0, nil, errBrokenSwitchboard
}
return id, conn, nil
return sb.conns[r], nil
}
// actively triggered by session.Close()
@ -148,16 +152,16 @@ func (sb *switchboard) closeAll() {
if !atomic.CompareAndSwapUint32(&sb.broken, 0, 1) {
return
}
sb.conns.Range(func(key, connI interface{}) bool {
conn := connI.(net.Conn)
sb.connsM.Lock()
for _, conn := range sb.conns {
conn.Close()
sb.conns.Delete(key)
return true
})
}
sb.conns = nil
sb.connsM.Unlock()
}
// deplex function costantly reads from a TCP connection
func (sb *switchboard) deplex(connId uint32, conn net.Conn) {
func (sb *switchboard) deplex(conn net.Conn) {
defer conn.Close()
buf := make([]byte, sb.session.connReceiveBufferSize)
for {
@ -166,8 +170,7 @@ func (sb *switchboard) deplex(connId uint32, conn net.Conn) {
sb.valve.AddRx(int64(n))
if err != nil {
log.Debugf("a connection for session %v has closed: %v", sb.session.id, err)
sb.conns.Delete(connId)
atomic.AddUint32(&sb.numConns, ^uint32(0))
sb.delConn(conn)
sb.session.SetTerminalMsg("a connection has dropped unexpectedly")
sb.session.passiveClose()
return

32
internal/multiplex/switchboard_test.go
View File

@ -14,14 +14,14 @@ func TestSwitchboard_Send(t *testing.T) {
sesh := MakeSession(0, seshConfig)
hole0 := connutil.Discard()
sesh.sb.addConn(hole0)
connId, _, err := sesh.sb.pickRandConn()
conn, err := sesh.sb.pickRandConn()
if err != nil {
t.Error("failed to get a random conn", err)
return
}
data := make([]byte, 1000)
rand.Read(data)
_, err = sesh.sb.send(data, &connId)
_, err = sesh.sb.send(data, &conn)
if err != nil {
t.Error(err)
return
@ -29,23 +29,23 @@ func TestSwitchboard_Send(t *testing.T) {
hole1 := connutil.Discard()
sesh.sb.addConn(hole1)
connId, _, err = sesh.sb.pickRandConn()
conn, err = sesh.sb.pickRandConn()
if err != nil {
t.Error("failed to get a random conn", err)
return
}
_, err = sesh.sb.send(data, &connId)
_, err = sesh.sb.send(data, &conn)
if err != nil {
t.Error(err)
return
}
connId, _, err = sesh.sb.pickRandConn()
conn, err = sesh.sb.pickRandConn()
if err != nil {
t.Error("failed to get a random conn", err)
return
}
_, err = sesh.sb.send(data, &connId)
_, err = sesh.sb.send(data, &conn)
if err != nil {
t.Error(err)
return
@ -71,7 +71,7 @@ func BenchmarkSwitchboard_Send(b *testing.B) {
seshConfig := SessionConfig{}
sesh := MakeSession(0, seshConfig)
sesh.sb.addConn(hole)
connId, _, err := sesh.sb.pickRandConn()
conn, err := sesh.sb.pickRandConn()
if err != nil {
b.Error("failed to get a random conn", err)
return
@ -81,7 +81,7 @@ func BenchmarkSwitchboard_Send(b *testing.B) {
b.SetBytes(int64(len(data)))
b.ResetTimer()
for i := 0; i < b.N; i++ {
sesh.sb.send(data, &connId)
sesh.sb.send(data, &conn)
}
}
@ -92,7 +92,7 @@ func TestSwitchboard_TxCredit(t *testing.T) {
sesh := MakeSession(0, seshConfig)
hole := connutil.Discard()
sesh.sb.addConn(hole)
connId, _, err := sesh.sb.pickRandConn()
conn, err := sesh.sb.pickRandConn()
if err != nil {
t.Error("failed to get a random conn", err)
return
@ -103,7 +103,7 @@ func TestSwitchboard_TxCredit(t *testing.T) {
t.Run("FIXED CONN MAPPING", func(t *testing.T) {
*sesh.sb.valve.(*LimitedValve).tx = 0
sesh.sb.strategy = FIXED_CONN_MAPPING
n, err := sesh.sb.send(data[:10], &connId)
n, err := sesh.sb.send(data[:10], &conn)
if err != nil {
t.Error(err)
return
@ -119,7 +119,7 @@ func TestSwitchboard_TxCredit(t *testing.T) {
t.Run("UNIFORM", func(t *testing.T) {
*sesh.sb.valve.(*LimitedValve).tx = 0
sesh.sb.strategy = UNIFORM_SPREAD
n, err := sesh.sb.send(data[:10], &connId)
n, err := sesh.sb.send(data[:10], &conn)
if err != nil {
t.Error(err)
return
@ -173,13 +173,17 @@ func TestSwitchboard_ConnsCount(t *testing.T) {
}
wg.Wait()
if sesh.sb.connsCount() != 1000 {
sesh.sb.connsM.RLock()
if len(sesh.sb.conns) != 1000 {
t.Error("connsCount incorrect")
}
sesh.sb.connsM.RUnlock()
sesh.sb.closeAll()
assert.Eventuallyf(t, func() bool {
return sesh.sb.connsCount() == 0
}, time.Second, 10*time.Millisecond, "connsCount incorrect: %v", sesh.sb.connsCount())
sesh.sb.connsM.RLock()
defer sesh.sb.connsM.RUnlock()
return len(sesh.sb.conns) == 0
}, time.Second, 10*time.Millisecond, "connsCount incorrect")
}