croc/src/tcp/tcp.go

179 lines
3.6 KiB
Go
Raw Normal View History

2018-09-23 21:34:29 +02:00
package tcp
import (
2018-09-26 16:41:07 +02:00
"bufio"
2018-09-23 21:34:29 +02:00
"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
2018-09-23 21:34:29 +02:00
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)
}
2018-09-23 22:22:57 +02:00
// TODO:
// delete old rooms
2018-09-23 21:34:29 +02:00
}
func run(port string) (err error) {
log.Debugf("starting TCP server on " + port)
2018-09-26 17:54:01 +02:00
rAddr, err := net.ResolveTCPAddr("tcp", "0.0.0.0:"+port)
2018-09-26 17:46:12 +02:00
if err != nil {
panic(err)
}
server, err := net.ListenTCP("tcp", rAddr)
2018-09-23 21:34:29 +02:00
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) {
2018-09-23 21:34:29 +02:00
// 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) {
2018-09-23 21:34:29 +02:00
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)
2018-09-26 16:41:07 +02:00
reader := bufio.NewReader(conn)
2018-09-23 21:34:29 +02:00
go func() {
for {
2018-09-26 18:11:36 +02:00
b := make([]byte, models.TCP_BUFFER_SIZE)
2018-09-26 16:41:07 +02:00
n, err := reader.Read(b)
2018-09-26 15:36:47 +02:00
if n > 0 {
2018-09-26 18:13:11 +02:00
c <- b[:n]
2018-09-26 18:12:38 +02:00
// res := make([]byte, n)
// // Copy the buffer so it doesn't get changed while read by the recipient.
// copy(res, b[:n])
// c <- res
2018-09-26 15:36:47 +02:00
}
2018-09-23 21:34:29 +02:00
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)
2018-09-26 16:41:07 +02:00
writer1 := bufio.NewWriter(conn1)
writer2 := bufio.NewWriter(conn2)
2018-09-23 21:34:29 +02:00
for {
select {
case b1 := <-chan1:
if b1 == nil {
return
}
2018-09-26 16:41:07 +02:00
writer2.Write(b1)
2018-09-23 21:34:29 +02:00
case b2 := <-chan2:
if b2 == nil {
return
}
2018-09-26 16:41:07 +02:00
writer1.Write(b2)
2018-09-23 21:34:29 +02:00
}
}
}