package tcp import ( "bufio" "net" "sync" "time" log "github.com/cihub/seelog" "github.com/pkg/errors" "github.com/schollz/croc/src/comm" "github.com/schollz/croc/src/logger" "github.com/schollz/croc/src/models" ) type roomInfo struct { receiver *comm.Comm opened time.Time } type roomMap struct { rooms map[string]roomInfo sync.Mutex } var rooms roomMap // Run starts a tcp listener, run async func Run(debugLevel, port string) { logger.SetLogLevel(debugLevel) rooms.Lock() rooms.rooms = make(map[string]roomInfo) rooms.Unlock() err := run(port) if err != nil { log.Error(err) } // TODO: // delete old rooms } func run(port string) (err error) { log.Debugf("starting TCP server on " + port) rAddr, err := net.ResolveTCPAddr("tcp", "0.0.0.0:"+port) if err != nil { panic(err) } server, err := net.ListenTCP("tcp", rAddr) if err != nil { return errors.Wrap(err, "Error listening on :"+port) } defer server.Close() // spawn a new goroutine whenever a client connects for { connection, err := server.Accept() if err != nil { return errors.Wrap(err, "problem accepting connection") } log.Debugf("client %s connected", connection.RemoteAddr().String()) go func(port string, connection net.Conn) { errCommunication := clientCommuncation(port, comm.New(connection)) if errCommunication != nil { log.Warnf("relay-%s: %s", connection.RemoteAddr().String(), errCommunication.Error()) } }(port, connection) } } func clientCommuncation(port string, c *comm.Comm) (err error) { // send ok to tell client they are connected err = c.Send("ok") if err != nil { return } // wait for client to tell me which room they want room, err := c.Receive() if err != nil { return } rooms.Lock() // first connection is always the receiver if _, ok := rooms.rooms[room]; !ok { rooms.rooms[room] = roomInfo{ receiver: c, opened: time.Now(), } rooms.Unlock() // tell the client that they got the room err = c.Send("recipient") if err != nil { return } return nil } receiver := rooms.rooms[room].receiver rooms.Unlock() // second connection is the sender, time to staple connections var wg sync.WaitGroup wg.Add(1) // start piping go func(com1, com2 *comm.Comm, wg *sync.WaitGroup) { log.Debug("starting pipes") pipe(com1.Connection(), com2.Connection()) wg.Done() log.Debug("done piping") }(c, receiver, &wg) // tell the sender everything is ready err = c.Send("sender") if err != nil { return } wg.Wait() // delete room rooms.Lock() log.Debugf("deleting room: %s", room) delete(rooms.rooms, room) rooms.Unlock() return nil } // chanFromConn creates a channel from a Conn object, and sends everything it // Read()s from the socket to the channel. func chanFromConn(conn net.Conn) chan []byte { c := make(chan []byte) reader := bufio.NewReader(conn) go func() { for { b := make([]byte, models.TCP_BUFFER_SIZE) n, err := reader.Read(b) if n > 0 { c <- b[:n] // res := make([]byte, n) // // Copy the buffer so it doesn't get changed while read by the recipient. // copy(res, b[:n]) // c <- res } if err != nil { c <- nil break } } }() return c } // pipe creates a full-duplex pipe between the two sockets and // transfers data from one to the other. func pipe(conn1 net.Conn, conn2 net.Conn) { chan1 := chanFromConn(conn1) chan2 := chanFromConn(conn2) writer1 := bufio.NewWriter(conn1) writer2 := bufio.NewWriter(conn2) for { select { case b1 := <-chan1: if b1 == nil { return } writer2.Write(b1) case b2 := <-chan2: if b2 == nil { return } writer1.Write(b2) } } }