diff --git a/api.go b/api.go index 9a69667..4c70c55 100644 --- a/api.go +++ b/api.go @@ -6,6 +6,7 @@ import ( "os" "time" + "github.com/googollee/go-socket.io" gh "github.com/gorilla/handlers" "github.com/gorilla/mux" "github.com/play-with-docker/play-with-docker/config" @@ -21,7 +22,14 @@ func main() { bypassCaptcha := len(os.Getenv("GOOGLE_RECAPTCHA_DISABLED")) > 0 - server := handlers.Broadcast.GetHandler() + server, err := socketio.NewServer(nil) + if err != nil { + log.Fatal(err) + } + server.On("connection", handlers.WS) + server.On("error", handlers.WSError) + + handlers.RegisterEvents(server) r := mux.NewRouter() corsRouter := mux.NewRouter() diff --git a/event/event.go b/event/event.go index 6a6541f..3216a88 100644 --- a/event/event.go +++ b/event/event.go @@ -2,15 +2,24 @@ package event type EventType string +func (e EventType) String() string { + return string(e) +} + const INSTANCE_VIEWPORT_RESIZE EventType = "instance viewport resize" const INSTANCE_DELETE EventType = "instance delete" const INSTANCE_NEW EventType = "instance new" +const INSTANCE_STATS EventType = "instance stats" +const INSTANCE_TERMINAL_OUT EventType = "instance terminal out" const SESSION_END EventType = "session end" const SESSION_READY EventType = "session ready" +const SESSION_BUILDER_OUT EventType = "session builder out" -type Handler func(args ...interface{}) +type Handler func(sessionId string, args ...interface{}) +type AnyHandler func(eventType EventType, sessionId string, args ...interface{}) type EventApi interface { - Emit(name EventType, args ...interface{}) + Emit(name EventType, sessionId string, args ...interface{}) On(name EventType, handler Handler) + OnAny(handler AnyHandler) } diff --git a/event/local_broker.go b/event/local_broker.go index da68dd4..2f33292 100644 --- a/event/local_broker.go +++ b/event/local_broker.go @@ -5,11 +5,12 @@ import "sync" type localBroker struct { sync.Mutex - handlers map[EventType][]Handler + handlers map[EventType][]Handler + anyHandlers []AnyHandler } func NewLocalBroker() *localBroker { - return &localBroker{handlers: map[EventType][]Handler{}} + return &localBroker{handlers: map[EventType][]Handler{}, anyHandlers: []AnyHandler{}} } func (b *localBroker) On(name EventType, handler Handler) { @@ -22,13 +23,25 @@ func (b *localBroker) On(name EventType, handler Handler) { b.handlers[name] = append(b.handlers[name], handler) } -func (b *localBroker) Emit(name EventType, args ...interface{}) { +func (b *localBroker) OnAny(handler AnyHandler) { b.Lock() defer b.Unlock() - if b.handlers[name] != nil { - for _, handler := range b.handlers[name] { - handler(args...) - } - } + b.anyHandlers = append(b.anyHandlers, handler) +} + +func (b *localBroker) Emit(name EventType, sessionId string, args ...interface{}) { + go func() { + b.Lock() + defer b.Unlock() + + for _, handler := range b.anyHandlers { + handler(name, sessionId, args...) + } + if b.handlers[name] != nil { + for _, handler := range b.handlers[name] { + handler(sessionId, args...) + } + } + }() } diff --git a/event/local_broker_test.go b/event/local_broker_test.go index 777a4b8..90939c8 100644 --- a/event/local_broker_test.go +++ b/event/local_broker_test.go @@ -7,25 +7,54 @@ import ( "github.com/stretchr/testify/assert" ) -func TestLocalBroker(t *testing.T) { +func TestLocalBroker_On(t *testing.T) { broker := NewLocalBroker() called := 0 + receivedSessionId := "" receivedArgs := []interface{}{} wg := sync.WaitGroup{} wg.Add(1) - broker.On(INSTANCE_NEW, func(args ...interface{}) { + broker.On(INSTANCE_NEW, func(sessionId string, args ...interface{}) { called++ + receivedSessionId = sessionId receivedArgs = args wg.Done() }) - broker.Emit(SESSION_READY) - broker.Emit(INSTANCE_NEW, "foo", "bar") + broker.Emit(SESSION_READY, "1") + broker.Emit(INSTANCE_NEW, "2", "foo", "bar") wg.Wait() assert.Equal(t, 1, called) + assert.Equal(t, "2", receivedSessionId) assert.Equal(t, []interface{}{"foo", "bar"}, receivedArgs) } + +func TestLocalBroker_OnAny(t *testing.T) { + broker := NewLocalBroker() + + var receivedEvent EventType + receivedSessionId := "" + receivedArgs := []interface{}{} + + wg := sync.WaitGroup{} + wg.Add(1) + + broker.OnAny(func(eventType EventType, sessionId string, args ...interface{}) { + receivedSessionId = sessionId + receivedArgs = args + receivedEvent = eventType + wg.Done() + }) + broker.Emit(SESSION_READY, "1") + + wg.Wait() + + var expectedArgs []interface{} + assert.Equal(t, SESSION_READY, receivedEvent) + assert.Equal(t, "1", receivedSessionId) + assert.Equal(t, expectedArgs, receivedArgs) +} diff --git a/handlers/bootstrap.go b/handlers/bootstrap.go index 9e8e90b..ffd17f0 100644 --- a/handlers/bootstrap.go +++ b/handlers/bootstrap.go @@ -5,14 +5,17 @@ import ( "os" "github.com/docker/docker/client" + "github.com/googollee/go-socket.io" "github.com/play-with-docker/play-with-docker/config" "github.com/play-with-docker/play-with-docker/docker" + "github.com/play-with-docker/play-with-docker/event" "github.com/play-with-docker/play-with-docker/pwd" "github.com/play-with-docker/play-with-docker/storage" ) var core pwd.PWDApi -var Broadcast pwd.BroadcastApi +var e event.EventApi +var ws *socketio.Server func Bootstrap() { c, err := client.NewEnvClient() @@ -22,18 +25,24 @@ func Bootstrap() { d := docker.NewDocker(c) - Broadcast, err = pwd.NewBroadcast(WS, WSError) - if err != nil { - log.Fatal(err) - } + e = event.NewLocalBroker() - t := pwd.NewScheduler(Broadcast, d) + t := pwd.NewScheduler(e, d) s, err := storage.NewFileStorage(config.SessionsFile) if err != nil && !os.IsNotExist(err) { log.Fatal("Error decoding sessions from disk ", err) } - core = pwd.NewPWD(d, t, Broadcast, s) + core = pwd.NewPWD(d, t, e, s) } + +func RegisterEvents(s *socketio.Server) { + ws = s + e.OnAny(broadcastEvent) +} + +func broadcastEvent(eventType event.EventType, sessionId string, args ...interface{}) { + ws.BroadcastTo(sessionId, eventType.String(), args...) +} diff --git a/pwd/broadcast.go b/pwd/broadcast.go deleted file mode 100644 index 1f54d08..0000000 --- a/pwd/broadcast.go +++ /dev/null @@ -1,34 +0,0 @@ -package pwd - -import ( - "net/http" - - "github.com/googollee/go-socket.io" -) - -type BroadcastApi interface { - BroadcastTo(sessionId, eventName string, args ...interface{}) - GetHandler() http.Handler -} - -type broadcast struct { - sio *socketio.Server -} - -func (b *broadcast) BroadcastTo(sessionId, eventName string, args ...interface{}) { - b.sio.BroadcastTo(sessionId, eventName, args...) -} - -func (b *broadcast) GetHandler() http.Handler { - return b.sio -} - -func NewBroadcast(connectionEvent, errorEvent interface{}) (*broadcast, error) { - server, err := socketio.NewServer(nil) - if err != nil { - return nil, err - } - server.On("connection", connectionEvent) - server.On("error", errorEvent) - return &broadcast{sio: server}, nil -} diff --git a/pwd/broadcast_mock_test.go b/pwd/broadcast_mock_test.go deleted file mode 100644 index 1d40cdc..0000000 --- a/pwd/broadcast_mock_test.go +++ /dev/null @@ -1,20 +0,0 @@ -package pwd - -import "net/http" - -type mockBroadcast struct { - broadcastTo func(sessionId, eventName string, args ...interface{}) - getHandler func() http.Handler -} - -func (m *mockBroadcast) BroadcastTo(sessionId, eventName string, args ...interface{}) { - if m.broadcastTo != nil { - m.broadcastTo(sessionId, eventName, args...) - } -} -func (m *mockBroadcast) GetHandler() http.Handler { - if m.getHandler != nil { - return m.getHandler() - } - return nil -} diff --git a/pwd/client.go b/pwd/client.go index 6fc1021..10660bd 100644 --- a/pwd/client.go +++ b/pwd/client.go @@ -4,6 +4,7 @@ import ( "log" "time" + "github.com/play-with-docker/play-with-docker/event" "github.com/play-with-docker/play-with-docker/pwd/types" ) @@ -41,7 +42,7 @@ func (p *pwd) ClientClose(client *types.Client) { func (p *pwd) notifyClientSmallestViewPort(session *types.Session) { vp := p.SessionGetSmallestViewPort(session) // Resize all terminals in the session - p.broadcast.BroadcastTo(session.Id, "viewport resize", vp.Cols, vp.Rows) + p.event.Emit(event.INSTANCE_VIEWPORT_RESIZE, session.Id, vp.Cols, vp.Rows) for _, instance := range session.Instances { err := p.InstanceResizeTerminal(instance, vp.Rows, vp.Cols) if err != nil { diff --git a/pwd/client_test.go b/pwd/client_test.go index 7ab9243..e682fdb 100644 --- a/pwd/client_test.go +++ b/pwd/client_test.go @@ -1,9 +1,11 @@ package pwd import ( + "sync" "testing" "time" + "github.com/play-with-docker/play-with-docker/event" "github.com/play-with-docker/play-with-docker/pwd/types" "github.com/stretchr/testify/assert" ) @@ -11,10 +13,10 @@ import ( func TestClientNew(t *testing.T) { docker := &mockDocker{} tasks := &mockTasks{} - broadcast := &mockBroadcast{} + e := event.NewLocalBroker() storage := &mockStorage{} - p := NewPWD(docker, tasks, broadcast, storage) + p := NewPWD(docker, tasks, e, storage) session, err := p.SessionNew(time.Hour, "", "", "") assert.Nil(t, err) @@ -26,33 +28,34 @@ func TestClientNew(t *testing.T) { } func TestClientResizeViewPort(t *testing.T) { + wg := sync.WaitGroup{} + wg.Add(1) docker := &mockDocker{} tasks := &mockTasks{} - broadcast := &mockBroadcast{} + e := event.NewLocalBroker() broadcastedSessionId := "" - broadcastedEventName := "" broadcastedArgs := []interface{}{} - broadcast.broadcastTo = func(sessionId, eventName string, args ...interface{}) { + e.On(event.INSTANCE_VIEWPORT_RESIZE, func(sessionId string, args ...interface{}) { broadcastedSessionId = sessionId - broadcastedEventName = eventName broadcastedArgs = args - } + wg.Done() + }) storage := &mockStorage{} - p := NewPWD(docker, tasks, broadcast, storage) + p := NewPWD(docker, tasks, e, storage) session, err := p.SessionNew(time.Hour, "", "", "") assert.Nil(t, err) client := p.ClientNew("foobar", session) p.ClientResizeViewPort(client, 80, 24) + wg.Wait() assert.Equal(t, types.ViewPort{Cols: 80, Rows: 24}, client.ViewPort) assert.Equal(t, session.Id, broadcastedSessionId) - assert.Equal(t, "viewport resize", broadcastedEventName) assert.Equal(t, uint(80), broadcastedArgs[0]) assert.Equal(t, uint(24), broadcastedArgs[1]) } diff --git a/pwd/instance.go b/pwd/instance.go index 5af7544..4c45c3d 100644 --- a/pwd/instance.go +++ b/pwd/instance.go @@ -11,6 +11,7 @@ import ( "github.com/play-with-docker/play-with-docker/config" "github.com/play-with-docker/play-with-docker/docker" + "github.com/play-with-docker/play-with-docker/event" "github.com/play-with-docker/play-with-docker/pwd/types" "golang.org/x/text/encoding" @@ -19,11 +20,11 @@ import ( type sessionWriter struct { sessionId string instanceName string - broadcast BroadcastApi + event event.EventApi } func (s *sessionWriter) Write(p []byte) (n int, err error) { - s.broadcast.BroadcastTo(s.sessionId, "terminal out", s.instanceName, string(p)) + s.event.Emit(event.INSTANCE_TERMINAL_OUT, s.sessionId, s.instanceName, string(p)) return len(p), nil } @@ -52,7 +53,7 @@ func (p *pwd) InstanceAttachTerminal(instance *types.Instance) error { } encoder := encoding.Replacement.NewEncoder() - sw := &sessionWriter{sessionId: instance.Session.Id, instanceName: instance.Name, broadcast: p.broadcast} + sw := &sessionWriter{sessionId: instance.Session.Id, instanceName: instance.Name, event: p.event} instance.Terminal = conn io.Copy(encoder.Writer(sw), conn) @@ -139,7 +140,7 @@ func (p *pwd) InstanceDelete(session *types.Session, instance *types.Instance) e return err } - p.broadcast.BroadcastTo(session.Id, "delete instance", instance.Name) + p.event.Emit(event.INSTANCE_DELETE, session.Id, instance.Name) delete(session.Instances, instance.Name) if err := p.storage.SessionPut(session); err != nil { @@ -238,7 +239,7 @@ func (p *pwd) InstanceNew(session *types.Session, conf InstanceConfig) (*types.I return nil, err } - p.broadcast.BroadcastTo(session.Id, "new instance", instance.Name, instance.IP, instance.Hostname) + p.event.Emit(event.INSTANCE_NEW, session.Id, instance.Name, instance.IP, instance.Hostname) p.setGauges() diff --git a/pwd/instance_test.go b/pwd/instance_test.go index cb82ac1..334ed80 100644 --- a/pwd/instance_test.go +++ b/pwd/instance_test.go @@ -8,6 +8,7 @@ import ( "github.com/play-with-docker/play-with-docker/config" "github.com/play-with-docker/play-with-docker/docker" + "github.com/play-with-docker/play-with-docker/event" "github.com/play-with-docker/play-with-docker/pwd/types" "github.com/stretchr/testify/assert" ) @@ -27,10 +28,10 @@ func TestInstanceResizeTerminal(t *testing.T) { } tasks := &mockTasks{} - broadcast := &mockBroadcast{} + e := event.NewLocalBroker() storage := &mockStorage{} - p := NewPWD(docker, tasks, broadcast, storage) + p := NewPWD(docker, tasks, e, storage) err := p.InstanceResizeTerminal(&types.Instance{Name: "foobar"}, 24, 80) @@ -49,10 +50,10 @@ func TestInstanceNew(t *testing.T) { } tasks := &mockTasks{} - broadcast := &mockBroadcast{} + e := event.NewLocalBroker() storage := &mockStorage{} - p := NewPWD(dock, tasks, broadcast, storage) + p := NewPWD(dock, tasks, e, storage) session, err := p.SessionNew(time.Hour, "", "", "") @@ -99,10 +100,10 @@ func TestInstanceNew_Concurrency(t *testing.T) { } tasks := &mockTasks{} - broadcast := &mockBroadcast{} + e := event.NewLocalBroker() storage := &mockStorage{} - p := NewPWD(dock, tasks, broadcast, storage) + p := NewPWD(dock, tasks, e, storage) session, err := p.SessionNew(time.Hour, "", "", "") @@ -140,10 +141,10 @@ func TestInstanceNew_WithNotAllowedImage(t *testing.T) { } tasks := &mockTasks{} - broadcast := &mockBroadcast{} + e := event.NewLocalBroker() storage := &mockStorage{} - p := NewPWD(dock, tasks, broadcast, storage) + p := NewPWD(dock, tasks, e, storage) session, err := p.SessionNew(time.Hour, "", "", "") @@ -188,10 +189,10 @@ func TestInstanceNew_WithCustomHostname(t *testing.T) { } tasks := &mockTasks{} - broadcast := &mockBroadcast{} + e := event.NewLocalBroker() storage := &mockStorage{} - p := NewPWD(dock, tasks, broadcast, storage) + p := NewPWD(dock, tasks, e, storage) session, err := p.SessionNew(time.Hour, "", "", "") @@ -230,10 +231,10 @@ func TestInstanceNew_WithCustomHostname(t *testing.T) { func TestInstanceAllowedImages(t *testing.T) { dock := &mockDocker{} tasks := &mockTasks{} - broadcast := &mockBroadcast{} + e := event.NewLocalBroker() storage := &mockStorage{} - p := NewPWD(dock, tasks, broadcast, storage) + p := NewPWD(dock, tasks, e, storage) expectedImages := []string{config.GetDindImageName(), "franela/dind:overlay2-dev", "franela/ucp:2.4.1"} diff --git a/pwd/pwd.go b/pwd/pwd.go index 1dcc8fa..a4c0496 100644 --- a/pwd/pwd.go +++ b/pwd/pwd.go @@ -5,6 +5,7 @@ import ( "time" "github.com/play-with-docker/play-with-docker/docker" + "github.com/play-with-docker/play-with-docker/event" "github.com/play-with-docker/play-with-docker/pwd/types" "github.com/play-with-docker/play-with-docker/storage" "github.com/prometheus/client_golang/prometheus" @@ -43,10 +44,10 @@ func init() { } type pwd struct { - docker docker.DockerApi - tasks SchedulerApi - broadcast BroadcastApi - storage storage.StorageApi + docker docker.DockerApi + tasks SchedulerApi + event event.EventApi + storage storage.StorageApi } type PWDApi interface { @@ -76,8 +77,8 @@ type PWDApi interface { ClientClose(client *types.Client) } -func NewPWD(d docker.DockerApi, t SchedulerApi, b BroadcastApi, s storage.StorageApi) *pwd { - return &pwd{docker: d, tasks: t, broadcast: b, storage: s} +func NewPWD(d docker.DockerApi, t SchedulerApi, e event.EventApi, s storage.StorageApi) *pwd { + return &pwd{docker: d, tasks: t, event: e, storage: s} } func (p *pwd) setGauges() { diff --git a/pwd/session.go b/pwd/session.go index ad36812..dc0ab7b 100644 --- a/pwd/session.go +++ b/pwd/session.go @@ -11,17 +11,18 @@ import ( "github.com/play-with-docker/play-with-docker/config" "github.com/play-with-docker/play-with-docker/docker" + "github.com/play-with-docker/play-with-docker/event" "github.com/play-with-docker/play-with-docker/pwd/types" "github.com/twinj/uuid" ) type sessionBuilderWriter struct { sessionId string - broadcast BroadcastApi + event event.EventApi } func (s *sessionBuilderWriter) Write(p []byte) (n int, err error) { - s.broadcast.BroadcastTo(s.sessionId, "session builder out", string(p)) + s.event.Emit(event.SESSION_BUILDER_OUT, s.sessionId, string(p)) return len(p), nil } @@ -87,8 +88,7 @@ func (p *pwd) SessionClose(s *types.Session) error { s.StopTicker() - p.broadcast.BroadcastTo(s.Id, "session end") - p.broadcast.BroadcastTo(s.Id, "disconnect") + p.event.Emit(event.SESSION_END, s.Id) log.Printf("Starting clean up of session [%s]\n", s.Id) for _, i := range s.Instances { err := p.InstanceDelete(s, i) @@ -146,7 +146,7 @@ func (p *pwd) SessionDeployStack(s *types.Session) error { } s.Ready = false - p.broadcast.BroadcastTo(s.Id, "session ready", false) + p.event.Emit(event.SESSION_READY, s.Id, false) i, err := p.InstanceNew(s, InstanceConfig{ImageName: s.ImageName, Host: s.Host}) if err != nil { log.Printf("Error creating instance for stack [%s]: %s\n", s.Stack, err) @@ -162,7 +162,7 @@ func (p *pwd) SessionDeployStack(s *types.Session) error { file := fmt.Sprintf("/var/run/pwd/uploads/%s", fileName) cmd := fmt.Sprintf("docker swarm init --advertise-addr eth0 && docker-compose -f %s pull && docker stack deploy -c %s %s", file, file, s.StackName) - w := sessionBuilderWriter{sessionId: s.Id, broadcast: p.broadcast} + w := sessionBuilderWriter{sessionId: s.Id, event: p.event} code, err := p.docker.ExecAttach(i.Name, []string{"sh", "-c", cmd}, &w) if err != nil { log.Printf("Error executing stack [%s]: %s\n", s.Stack, err) @@ -171,7 +171,7 @@ func (p *pwd) SessionDeployStack(s *types.Session) error { log.Printf("Stack execution finished with code %d\n", code) s.Ready = true - p.broadcast.BroadcastTo(s.Id, "session ready", true) + p.event.Emit(event.SESSION_READY, s.Id, true) if err := p.storage.SessionPut(s); err != nil { return err } diff --git a/pwd/session_test.go b/pwd/session_test.go index bef5a1d..4e9ac53 100644 --- a/pwd/session_test.go +++ b/pwd/session_test.go @@ -7,6 +7,7 @@ import ( "github.com/play-with-docker/play-with-docker/config" "github.com/play-with-docker/play-with-docker/docker" + "github.com/play-with-docker/play-with-docker/event" "github.com/play-with-docker/play-with-docker/pwd/types" "github.com/stretchr/testify/assert" ) @@ -36,14 +37,14 @@ func TestSessionNew(t *testing.T) { scheduledSession = s } - broadcast := &mockBroadcast{} + ev := event.NewLocalBroker() storage := &mockStorage{} storage.sessionPut = func(s *types.Session) error { saveCalled = true return nil } - p := NewPWD(docker, tasks, broadcast, storage) + p := NewPWD(docker, tasks, ev, storage) before := time.Now() @@ -166,10 +167,10 @@ func TestSessionSetup(t *testing.T) { return nil, nil } tasks := &mockTasks{} - broadcast := &mockBroadcast{} + ev := event.NewLocalBroker() storage := &mockStorage{} - p := NewPWD(dock, tasks, broadcast, storage) + p := NewPWD(dock, tasks, ev, storage) s, e := p.SessionNew(time.Hour, "", "", "") assert.Nil(t, e) diff --git a/pwd/tasks.go b/pwd/tasks.go index adc52ce..27d61c2 100644 --- a/pwd/tasks.go +++ b/pwd/tasks.go @@ -15,6 +15,7 @@ import ( "github.com/docker/docker/client" "github.com/docker/go-connections/tlsconfig" "github.com/play-with-docker/play-with-docker/docker" + "github.com/play-with-docker/play-with-docker/event" "github.com/play-with-docker/play-with-docker/pwd/types" ) @@ -28,7 +29,7 @@ type SchedulerApi interface { } type scheduler struct { - broadcast BroadcastApi + event event.EventApi periodicTasks []periodicTask } @@ -102,7 +103,7 @@ func (sch *scheduler) Schedule(s *types.Session) { sort.Sort(ins.Ports) ins.CleanUsedPorts() - sch.broadcast.BroadcastTo(ins.Session.Id, "instance stats", ins.Name, ins.Mem, ins.Cpu, ins.IsManager, ins.Ports) + sch.event.Emit(event.INSTANCE_STATS, ins.Session.Id, ins.Name, ins.Mem, ins.Cpu, ins.IsManager, ins.Ports) } } }() @@ -111,8 +112,8 @@ func (sch *scheduler) Schedule(s *types.Session) { func (sch *scheduler) Unschedule(s *types.Session) { } -func NewScheduler(b BroadcastApi, d docker.DockerApi) *scheduler { - s := &scheduler{broadcast: b} +func NewScheduler(e event.EventApi, d docker.DockerApi) *scheduler { + s := &scheduler{event: e} s.periodicTasks = []periodicTask{&collectStatsTask{docker: d}, &checkSwarmStatusTask{}, &checkUsedPortsTask{}, &checkSwarmUsedPortsTask{}} return s } diff --git a/router/l2/l2.go b/router/l2/l2.go index 4055294..1cce468 100644 --- a/router/l2/l2.go +++ b/router/l2/l2.go @@ -131,6 +131,6 @@ func main() { go monitorNetworks() r := router.NewRouter(director, config.SSHKeyPath) - r.Listen(":443", ":53", ":22") + r.ListenAndWait(":443", ":53", ":22") defer r.Close() } diff --git a/router/router.go b/router/router.go index 04a3094..4948370 100644 --- a/router/router.go +++ b/router/router.go @@ -33,6 +33,56 @@ type proxyRouter struct { } func (r *proxyRouter) Listen(httpAddr, dnsAddr, sshAddr string) { + l, err := net.Listen("tcp", httpAddr) + if err != nil { + log.Fatal(err) + } + r.httpListener = l + go func() { + for !r.closed { + conn, err := r.httpListener.Accept() + if err != nil { + continue + } + go r.handleConnection(conn) + } + }() + + dnsMux := dns.NewServeMux() + dnsMux.HandleFunc(".", r.dnsRequest) + r.udpDnsServer = &dns.Server{Addr: dnsAddr, Net: "udp", Handler: dnsMux} + r.tcpDnsServer = &dns.Server{Addr: dnsAddr, Net: "tcp", Handler: dnsMux} + + wg := sync.WaitGroup{} + wg.Add(2) + + r.udpDnsServer.NotifyStartedFunc = func() { + wg.Done() + } + r.tcpDnsServer.NotifyStartedFunc = func() { + wg.Done() + } + go r.udpDnsServer.ListenAndServe() + go r.tcpDnsServer.ListenAndServe() + wg.Wait() + + lssh, err := net.Listen("tcp", sshAddr) + if err != nil { + log.Fatal("failed to listen for connection: ", err) + } + r.sshListener = lssh + go func() { + for { + nConn, err := lssh.Accept() + if err != nil { + log.Fatal("failed to accept incoming connection: ", err) + } + + go r.sshHandle(nConn) + } + }() +} +func (r *proxyRouter) ListenAndWait(httpAddr, dnsAddr, sshAddr string) { listenWG := sync.WaitGroup{} l, err := net.Listen("tcp", httpAddr)