croc/src/croc/sending.go

package croc

import (
	"errors"
	"fmt"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"time"

	log "github.com/cihub/seelog"
	"github.com/gorilla/websocket"
	"github.com/schollz/croc/src/recipient"
	"github.com/schollz/croc/src/relay"
	"github.com/schollz/croc/src/sender"
	"github.com/schollz/peerdiscovery"
)

// Send the file
func (c *Croc) Send(fname, codephrase string) (err error) {
	log.Debugf("sending %s", fname)
	errChan := make(chan error)

	// normally attempt two connections
	waitingFor := 2

	// use public relay
	if !c.LocalOnly {
		go func() {
			// atttempt to connect to public relay
			errChan <- c.sendReceive(c.Address, c.AddressWebsocketPort, c.AddressTCPPort, fname, codephrase, true, false)
		}()
	} else {
		waitingFor = 1
	}

	// use local relay
	if !c.NoLocal {
		go func() {
			// start own relay and connect to it
			go relay.Run(c.RelayWebsocketPort, "")
			time.Sleep(250 * time.Millisecond) // race condition here, but this should work most of the time :(

			// broadcast for peer discovery
			go func() {
				log.Debug("listening for local croc relay...")
				go peerdiscovery.Discover(peerdiscovery.Settings{
					Limit:     1,
					TimeLimit: 600 * time.Second,
					Delay:     50 * time.Millisecond,
					Payload:   []byte(c.RelayWebsocketPort + "-" + c.RelayTCPPort),
				})
			}()

			// connect to own relay
			errChan <- c.sendReceive("localhost", c.RelayWebsocketPort, c.RelayTCPPort, fname, codephrase, true, true)
		}()
	} else {
		waitingFor = 1
	}

	err = <-errChan
	if err == nil || waitingFor == 1 {
		log.Debug("returning")
		return
	}
	log.Debug(err)
	return <-errChan
}

// Receive the file
func (c *Croc) Receive(codephrase string) (err error) {
	log.Debug("receiving")

	// use local relay first
	if !c.NoLocal {
		// try to discovery codephrase and server through peer network
		discovered, errDiscover := peerdiscovery.Discover(peerdiscovery.Settings{
			Limit:     1,
			TimeLimit: 300 * time.Millisecond,
			Delay:     50 * time.Millisecond,
			Payload:   []byte("checking"),
		})
		if errDiscover != nil {
			log.Debug(errDiscover)
		}
		if len(discovered) > 0 {
			log.Debugf("discovered %s:%s", discovered[0].Address, discovered[0].Payload)
			// see if we can actually connect to it
			timeout := time.Duration(200 * time.Millisecond)
			client := http.Client{
				Timeout: timeout,
			}
			resp, err := client.Get(fmt.Sprintf("http://%s:%s/", discovered[0].Address, discovered[0].Payload))
			if err == nil {
				if resp.StatusCode == http.StatusOK {
					// we connected, so use this
					ports := strings.Split(string(discovered[0].Payload), "-")
					if len(ports) != 2 {
						return errors.New("bad payload")
					}
					return c.sendReceive(discovered[0].Address, ports[0], ports[1], "", codephrase, false, true)
				}
			} else {
				log.Debugf("could not connect: %s", err.Error())
			}
		} else {
			log.Debug("discovered no peers")
		}
	}

	// use public relay
	if !c.LocalOnly {
		log.Debug("using public relay")
		return c.sendReceive(c.Address, c.AddressWebsocketPort, c.AddressTCPPort, "", codephrase, false, false)
	}

	return errors.New("must use local or public relay")
}

func (c *Croc) sendReceive(address, websocketPort, tcpPort, fname, codephrase string, isSender bool, isLocal bool) (err error) {
	defer log.Flush()
	if len(codephrase) < 4 {
		return fmt.Errorf("codephrase is too short")
	}

	// allow interrupts
	interrupt := make(chan os.Signal, 1)
	signal.Notify(interrupt, os.Interrupt)

	done := make(chan struct{})
	// connect to server
	log.Debugf("connecting to %s", address+"/ws?room="+codephrase[:3])
	if len(websocketPort) > 0 {
		address += ":" + websocketPort
	}
	sock, _, err := websocket.DefaultDialer.Dial("ws://"+address+"/ws?room="+codephrase[:3], nil)
	if err != nil {
		return
	}
	defer sock.Close()

	// tell the websockets we are connected
	err = sock.WriteMessage(websocket.BinaryMessage, []byte("connected"))
	if err != nil {
		return err
	}

	if isSender {
		// start peerdiscovery relay server
		go sender.Send(isLocal, done, sock, fname, codephrase, c.UseCompression, c.UseEncryption)
	} else {
		go recipient.Receive(isLocal, done, sock, codephrase, c.NoRecipientPrompt, c.Stdout)
	}

	for {
		select {
		case <-done:
			return nil
		case <-interrupt:
			log.Debug("interrupt")
			err = sock.WriteMessage(websocket.TextMessage, []byte("interrupt"))
			if err != nil {
				return err
			}
			time.Sleep(50 * time.Millisecond)

			// Cleanly close the connection by sending a close message and then
			// waiting (with timeout) for the server to close the connection.
			err := sock.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
			if err != nil {
				log.Debug("write close:", err)
				return nil
			}
			select {
			case <-done:
			case <-time.After(time.Second):
			}
			return nil
		}
	}
}

// Relay will start a relay on the specified port
func (c *Croc) Relay() (err error) {
	return relay.Run(c.RelayWebsocketPort, c.RelayTCPPort)
}