Robot
/
croc
mirror of https://github.com/schollz/croc.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

relay small refactor

This commit is contained in:
callmefever 2017-10-24 18:24:26 +03:00
parent 3fd9a345c4
commit 6114a66d18
1 changed files with 14 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
relay.go
View File

@ -43,15 +43,18 @@ type Relay struct {
} }
func NewRelay(flags *Flags) *Relay { func NewRelay(flags *Flags) *Relay {
r := new(Relay) r := &Relay{
r.Debug = flags.Debug Debug: flags.Debug,
r.NumberOfConnections = MAX_NUMBER_THREADS NumberOfConnections: MAX_NUMBER_THREADS,
}
log.SetFormatter(&log.TextFormatter{}) log.SetFormatter(&log.TextFormatter{})
if r.Debug { if r.Debug {
log.SetLevel(log.DebugLevel) log.SetLevel(log.DebugLevel)
} else { } else {
log.SetLevel(log.WarnLevel) log.SetLevel(log.WarnLevel)
} }
return r return r
} }
@ -85,8 +88,8 @@ func (r *Relay) listenerThread(id int, wg *sync.WaitGroup) {
}) })
defer wg.Done() defer wg.Done()
err := r.listener(id)
if err != nil { if err := r.listener(id); err != nil {
logger.Error(err) logger.Error(err)
} }
} }
@ -94,11 +97,11 @@ func (r *Relay) listenerThread(id int, wg *sync.WaitGroup) {
func (r *Relay) listener(id int) (err error) { func (r *Relay) listener(id int) (err error) {
port := strconv.Itoa(27001 + id) port := strconv.Itoa(27001 + id)
logger := log.WithFields(log.Fields{ logger := log.WithFields(log.Fields{
"function": "listener" + ":" + port, "function": "listener:" + port,
}) })
server, err := net.Listen("tcp", "0.0.0.0:"+port) server, err := net.Listen("tcp", "0.0.0.0:"+port)
if err != nil { if err != nil {
return errors.Wrap(err, "Error listening on "+":"+port) return errors.Wrap(err, "Error listening on :"+port)
} }
defer server.Close() defer server.Close()
logger.Debug("waiting for connections") logger.Debug("waiting for connections")
@ -129,7 +132,8 @@ func (r *Relay) clientCommuncation(id int, connection net.Conn) {
connectionType, codePhrase, metaData := m[0], m[1], m[2] connectionType, codePhrase, metaData := m[0], m[1], m[2]
key := codePhrase + "-" + strconv.Itoa(id) key := codePhrase + "-" + strconv.Itoa(id)
if connectionType == "s" { // sender connection switch connectionType {
case "s": // sender connection
if r.connections.IsSenderConnected(key) { if r.connections.IsSenderConnected(key) {
sendMessage("no", connection) sendMessage("no", connection)
return return
@ -179,8 +183,7 @@ func (r *Relay) clientCommuncation(id int, connection net.Conn) {
delete(r.connections.potentialReceivers, key) delete(r.connections.potentialReceivers, key)
r.connections.Unlock() r.connections.Unlock()
logger.Debug("deleted sender and receiver") logger.Debug("deleted sender and receiver")
} else if connectionType == "r" || connectionType == "c" { case "r", "c": // receiver
//receiver
if r.connections.IsPotentialReceiverConnected(key) { if r.connections.IsPotentialReceiverConnected(key) {
sendMessage("no", connection) sendMessage("no", connection)
return return
@ -226,10 +229,9 @@ func (r *Relay) clientCommuncation(id int, connection net.Conn) {
r.connections.receiver[key] = connection r.connections.receiver[key] = connection
r.connections.Unlock() r.connections.Unlock()
} }
} else { default:
logger.Debugf("Got unknown protocol: '%s'", connectionType) logger.Debugf("Got unknown protocol: '%s'", connectionType)
} }
return
} }
func sendMessage(message string, connection net.Conn) { func sendMessage(message string, connection net.Conn) {