// handle multiplex-ed connection
func handleMux(conn io.ReadWriteCloser, target string) {
// stream multiplex
var mux *yamux.Session
config := &yamux.Config{
AcceptBacklog: 256,
EnableKeepAlive: true,
KeepAliveInterval: 30 * time.Second,
ConnectionWriteTimeout: 30 * time.Second,
MaxStreamWindowSize: 16777216,
LogOutput: os.Stderr,
}
m, err := yamux.Server(conn, config)
if err != nil {
log.Println(err)
return
}
mux = m
defer mux.Close()
for {
p1, err := mux.Accept()
if err != nil {
log.Println(err)
return
}
p2, err := net.DialTimeout("tcp", target, 5*time.Second)
if err != nil {
log.Println(err)
return
}
go handleClient(p1, p2)
}
}
func handleCreateTunnel(session *yamux.Session, port int) (url string, err error) {
// TODO: Attempt to bind to local port
addr := fmt.Sprintf("127.0.0.1:%v", port)
log.WithFields(logrus.Fields{
"server": addr,
}).Info("attempting to connect to server")
local, err := net.Dial("tcp", addr)
if err != nil {
return "", errors.New("Could not establish connection to host port.")
}
log.WithFields(logrus.Fields{
"server": addr,
}).Info("connected to server")
// Establish stream
stream, err := session.Open()
if err != nil {
return "", errors.New("Could not create stream.")
}
// TODO: send control message
log.WithFields(logrus.Fields{
"server": addr,
}).Info("sending TunnelRequest")
req := &msg.TunnelRequest{
Port: port,
}
if err := msg.WriteMsg(stream, req); err != nil {
return "", err
}
// TODO: wait for ack w/ url
log.WithFields(logrus.Fields{
"server": addr,
}).Info("waiting for TunnelReply")
var res msg.TunnelReply
if err := msg.ReadMsgInto(stream, &res); err != nil {
return "", err
}
log.WithFields(logrus.Fields{
"server": addr,
"port": port,
"uri": res.URI,
}).Info("received TunnelReply")
url = res.URI
// TODO: Update global state for later control messages
// TODO: launch goroutine to proxy data
go proxy.Proxy(stream, local)
return url, nil
}
func acceptStreams(logger log.Logger, session *yamux.Session, streamCh chan net.Conn) grim.TaskFunc {
return func(ctx context.Context) {
defer close(streamCh)
for {
select {
case <-ctx.Done():
return
default:
stream, err := session.Accept()
if err != nil {
if err != io.EOF {
logger.Error("multiplex conn accept failed", "err", err)
}
return
}
streamCh <- stream
}
}
}
}