mirror of
https://github.com/cbeuw/Cloak.git
synced 2024-10-23 15:46:25 +00:00
119 lines
2.6 KiB
Go
119 lines
2.6 KiB
Go
|
package multiplex
|
||
|
|
||
|
import (
|
||
|
"net"
|
||
|
"sort"
|
||
|
)
|
||
|
|
||
|
const (
|
||
|
sentNotifyBacklog = 1024
|
||
|
dispatchBacklog = 10240
|
||
|
)
|
||
|
|
||
|
type switchboard struct {
|
||
|
session *Session
|
||
|
|
||
|
ces []*connEnclave
|
||
|
|
||
|
// For telling dispatcher how many bytes have been sent after Connection.send.
|
||
|
sentNotifyCh chan *sentNotifier
|
||
|
dispatCh chan []byte
|
||
|
}
|
||
|
|
||
|
// Some data comes from a Stream to be sent through one of the many
|
||
|
// remoteConn, but which remoteConn should we use to send the data?
|
||
|
//
|
||
|
// In this case, we pick the remoteConn that has about the smallest sendQueue.
|
||
|
// Though "smallest" is not guaranteed because it doesn't has to be
|
||
|
type connEnclave struct {
|
||
|
sb *switchboard
|
||
|
remoteConn net.Conn
|
||
|
sendQueue int
|
||
|
}
|
||
|
|
||
|
type byQ []*connEnclave
|
||
|
|
||
|
func (a byQ) Len() int {
|
||
|
return len(a)
|
||
|
}
|
||
|
func (a byQ) Swap(i, j int) {
|
||
|
a[i], a[j] = a[j], a[i]
|
||
|
}
|
||
|
func (a byQ) Less(i, j int) bool {
|
||
|
return a[i].sendQueue < a[j].sendQueue
|
||
|
}
|
||
|
|
||
|
func makeSwitchboard(conns []net.Conn, sesh *Session) *switchboard {
|
||
|
sb := &switchboard{
|
||
|
session: sesh,
|
||
|
ces: []*connEnclave{},
|
||
|
sentNotifyCh: make(chan *sentNotifier, sentNotifyBacklog),
|
||
|
dispatCh: make(chan []byte, dispatchBacklog),
|
||
|
}
|
||
|
for _, c := range conns {
|
||
|
ce := &connEnclave{
|
||
|
sb: sb,
|
||
|
remoteConn: c,
|
||
|
sendQueue: 0,
|
||
|
}
|
||
|
sb.ces = append(sb.ces, ce)
|
||
|
}
|
||
|
|
||
|
return sb
|
||
|
}
|
||
|
|
||
|
func (sb *switchboard) run() {
|
||
|
go startDispatcher()
|
||
|
go startDeplexer()
|
||
|
}
|
||
|
|
||
|
// Everytime after a remoteConn sends something, it constructs this struct
|
||
|
// Which is sent back to dispatch() through sentNotifyCh to tell dispatch
|
||
|
// how many bytes it has sent
|
||
|
type sentNotifier struct {
|
||
|
ce *connEnclave
|
||
|
sent int
|
||
|
}
|
||
|
|
||
|
func (ce *connEnclave) send(data []byte) {
|
||
|
// TODO: error handling
|
||
|
n, _ := ce.remoteConn.Write(data)
|
||
|
sn := &sentNotifier{
|
||
|
ce,
|
||
|
n,
|
||
|
}
|
||
|
ce.sb.sentNotifyCh <- sn
|
||
|
}
|
||
|
|
||
|
// Dispatcher sends data coming from a stream to a remote connection
|
||
|
func (sb *switchboard) startDispatcher() {
|
||
|
for {
|
||
|
select {
|
||
|
// dispatCh receives data from stream.Write
|
||
|
case data := <-sb.dispatCh:
|
||
|
go sb.ces[0].send(data)
|
||
|
sb.ces[0].sendQueue += len(data)
|
||
|
case notified := <-sb.sentNotifyCh:
|
||
|
notified.ce.sendQueue -= notified.sent
|
||
|
sort.Sort(byQ(sb.ces))
|
||
|
}
|
||
|
}
|
||
|
}
|
||
|
|
||
|
// Deplexer sends data coming from a remote connection to a stream
|
||
|
func (sb *switchboard) startDeplexer() {
|
||
|
for _, ce := range sb.ces {
|
||
|
go func() {
|
||
|
buf := make([]byte, 20480)
|
||
|
for {
|
||
|
sb.session.obfsedReader(ce.remoteConn, buf)
|
||
|
frame := sb.session.deobfs(buf)
|
||
|
if !sb.session.isStream(frame.StreamID) {
|
||
|
sb.session.addStream(frame.StreamID)
|
||
|
}
|
||
|
sb.session.getStream(frame.StreamID).recvNewFrame(frame)
|
||
|
}
|
||
|
}()
|
||
|
}
|
||
|
}
|