mirror of
https://github.com/cbeuw/Cloak.git
synced 2024-10-23 15:46:25 +00:00
178 lines
4.4 KiB
Go
178 lines
4.4 KiB
Go
package multiplex
|
|
|
|
import (
|
|
"log"
|
|
"net"
|
|
"sort"
|
|
)
|
|
|
|
const (
|
|
sentNotifyBacklog = 1024
|
|
dispatchBacklog = 10240
|
|
newConnBacklog = 8
|
|
)
|
|
|
|
// switchboard is responsible for keeping the reference of TLS connections between client and server
|
|
type switchboard struct {
|
|
session *Session
|
|
|
|
ces []*connEnclave
|
|
|
|
// For telling dispatcher how many bytes have been sent after Connection.send.
|
|
sentNotifyCh chan *sentNotifier
|
|
// dispatCh is used by streams to send new data to remote
|
|
dispatCh chan []byte
|
|
newConnCh chan net.Conn
|
|
closingCECh chan *connEnclave
|
|
die chan struct{}
|
|
closing bool
|
|
}
|
|
|
|
// 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
|
|
}
|
|
|
|
// It takes at least 1 conn to start a switchboard
|
|
func makeSwitchboard(conn net.Conn, sesh *Session) *switchboard {
|
|
sb := &switchboard{
|
|
session: sesh,
|
|
ces: []*connEnclave{},
|
|
sentNotifyCh: make(chan *sentNotifier, sentNotifyBacklog),
|
|
dispatCh: make(chan []byte, dispatchBacklog),
|
|
newConnCh: make(chan net.Conn, newConnBacklog),
|
|
closingCECh: make(chan *connEnclave, 5),
|
|
die: make(chan struct{}),
|
|
}
|
|
ce := &connEnclave{
|
|
sb: sb,
|
|
remoteConn: conn,
|
|
sendQueue: 0,
|
|
}
|
|
sb.ces = append(sb.ces, ce)
|
|
go sb.deplex(ce)
|
|
|
|
go sb.dispatch()
|
|
return sb
|
|
}
|
|
|
|
// 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, err := ce.remoteConn.Write(data)
|
|
if err != nil {
|
|
ce.sb.closingCECh <- ce
|
|
log.Println(err)
|
|
}
|
|
|
|
sn := &sentNotifier{
|
|
ce,
|
|
n,
|
|
}
|
|
ce.sb.sentNotifyCh <- sn
|
|
|
|
}
|
|
|
|
// Dispatcher sends data coming from a stream to a remote connection
|
|
// I used channels here because I didn't want to use mutex
|
|
func (sb *switchboard) dispatch() {
|
|
var dying bool
|
|
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))
|
|
case conn := <-sb.newConnCh:
|
|
log.Println("newConn")
|
|
newCe := &connEnclave{
|
|
sb: sb,
|
|
remoteConn: conn,
|
|
sendQueue: 0,
|
|
}
|
|
sb.ces = append(sb.ces, newCe)
|
|
go sb.deplex(newCe)
|
|
case closing := <-sb.closingCECh:
|
|
log.Println("Closing conn")
|
|
for i, ce := range sb.ces {
|
|
if closing == ce {
|
|
sb.ces = append(sb.ces[:i], sb.ces[i+1:]...)
|
|
break
|
|
}
|
|
}
|
|
if len(sb.ces) == 0 && !dying {
|
|
sb.session.Close()
|
|
}
|
|
case <-sb.die:
|
|
dying = true
|
|
for _, ce := range sb.ces {
|
|
ce.remoteConn.Close()
|
|
}
|
|
return
|
|
}
|
|
}
|
|
}
|
|
|
|
// deplex function costantly reads from a TCP connection
|
|
// it is responsible to act in response to the deobfsed header
|
|
// i.e. should a new stream be added? which existing stream should be closed?
|
|
func (sb *switchboard) deplex(ce *connEnclave) {
|
|
var highestStream uint32
|
|
buf := make([]byte, 20480)
|
|
for {
|
|
i, err := sb.session.obfsedReader(ce.remoteConn, buf)
|
|
if err != nil {
|
|
log.Println(err)
|
|
go ce.remoteConn.Close()
|
|
sb.closingCECh <- ce
|
|
return
|
|
}
|
|
frame := sb.session.deobfs(buf[:i])
|
|
if closing := sb.session.getStream(frame.ClosingStreamID); closing != nil {
|
|
log.Printf("HeaderClosing: %v\n", frame.ClosingStreamID)
|
|
closing.Close()
|
|
}
|
|
|
|
var stream *Stream
|
|
// If we want to open a new stream, we need to make sure that the newStreamID is indeed new
|
|
// i.e. it is not a stream that existed before but has been closed
|
|
// we don't allow streamID reuse.
|
|
// So here we do a check that the new stream has a higher ID than the highest ID we have got
|
|
if stream = sb.session.getStream(frame.StreamID); highestStream < frame.StreamID && stream == nil {
|
|
stream = sb.session.addStream(frame.StreamID)
|
|
highestStream = frame.StreamID
|
|
}
|
|
if stream != nil {
|
|
stream.newFrameCh <- frame
|
|
}
|
|
}
|
|
}
|