don't wait for ok

This commit is contained in:
Zack Scholl 2018-09-22 07:23:10 -07:00
parent 54b22cd732
commit 742a935ea9
3 changed files with 99 additions and 86 deletions

View File

@ -11,7 +11,6 @@ import (
"github.com/briandowns/spinner" "github.com/briandowns/spinner"
humanize "github.com/dustin/go-humanize" humanize "github.com/dustin/go-humanize"
"github.com/schollz/croc/src/zipper"
log "github.com/cihub/seelog" log "github.com/cihub/seelog"
"github.com/gorilla/websocket" "github.com/gorilla/websocket"
@ -20,6 +19,7 @@ import (
"github.com/schollz/croc/src/logger" "github.com/schollz/croc/src/logger"
"github.com/schollz/croc/src/models" "github.com/schollz/croc/src/models"
"github.com/schollz/croc/src/utils" "github.com/schollz/croc/src/utils"
"github.com/schollz/croc/src/zipper"
"github.com/schollz/pake" "github.com/schollz/pake"
"github.com/schollz/progressbar/v2" "github.com/schollz/progressbar/v2"
"github.com/tscholl2/siec" "github.com/tscholl2/siec"
@ -44,6 +44,7 @@ func receive(c *websocket.Conn, codephrase string, noPrompt bool) (err error) {
var fstats models.FileStats var fstats models.FileStats
var sessionKey []byte var sessionKey []byte
var transferTime time.Duration var transferTime time.Duration
var hash256 []byte
// start a spinner // start a spinner
spin := spinner.New(spinner.CharSets[9], 100*time.Millisecond) spin := spinner.New(spinner.CharSets[9], 100*time.Millisecond)
@ -146,90 +147,97 @@ func receive(c *websocket.Conn, codephrase string, noPrompt bool) (err error) {
if err != nil { if err != nil {
return err return err
} }
if messageType == websocket.PongMessage || messageType == websocket.PingMessage { if messageType != websocket.BinaryMessage {
continue continue
} }
if messageType == websocket.BinaryMessage {
// tell the sender that we recieved this packet
c.WriteMessage(websocket.BinaryMessage, []byte("ok"))
// do decryption // // tell the sender that we recieved this packet
var enc crypt.Encryption // c.WriteMessage(websocket.BinaryMessage, []byte("ok"))
err = json.Unmarshal(message, &enc)
if err != nil {
return err
}
decrypted, err := enc.Decrypt(sessionKey, true)
if err != nil {
return err
}
// do decompression // do decryption
decompressed := compress.Decompress(decrypted) var enc crypt.Encryption
// decompressed := decrypted err = json.Unmarshal(message, &enc)
if err != nil {
return err
}
decrypted, err := enc.Decrypt(sessionKey, true)
if err != nil {
return err
}
// write to file // do decompression
n, err := f.Write(decompressed) decompressed := compress.Decompress(decrypted)
if err != nil { // decompressed := decrypted
return err
}
// update the bytes written
bytesWritten += n
// update the progress bar
bar.Add(n)
} else {
// we are finished
transferTime = time.Since(startTime)
// close file // write to file
err = f.Close() n, err := f.Write(decompressed)
if err != nil { if err != nil {
return err return err
} }
// update the bytes written
bytesWritten += n
// update the progress bar
bar.Add(n)
// finish bar if int64(bytesWritten) == fstats.Size {
bar.Finish() break
// check hash
hash256, err := utils.HashFile(fstats.SentName)
if err != nil {
log.Error(err)
return err
}
// check success hash(myfile) == hash(theirfile)
log.Debugf("got hash: %x", message)
if bytes.Equal(hash256, message) {
c.WriteMessage(websocket.BinaryMessage, []byte("ok"))
// open directory
if fstats.IsDir {
err = zipper.UnzipFile(fstats.SentName, ".")
if DebugLevel != "debug" {
os.Remove(fstats.SentName)
}
} else {
err = nil
}
if err == nil {
transferRate := float64(fstats.Size) / 1000000.0 / transferTime.Seconds()
transferType := "MB/s"
if transferRate < 1 {
transferRate = float64(fstats.Size) / 1000.0 / transferTime.Seconds()
transferType = "kB/s"
}
folderOrFile := "file"
if fstats.IsDir {
folderOrFile = "folder"
}
fmt.Fprintf(os.Stderr, "\nReceived %s written to %s (%2.1f %s)\n", folderOrFile, fstats.Name, transferRate, transferType)
}
return err
} else {
c.WriteMessage(websocket.BinaryMessage, []byte("not"))
return errors.New("file corrupted")
}
} }
} }
c.WriteMessage(websocket.BinaryMessage, []byte("done"))
// we are finished
transferTime = time.Since(startTime)
// close file
err = f.Close()
if err != nil {
return err
}
// finish bar
bar.Finish()
// check hash
hash256, err = utils.HashFile(fstats.SentName)
if err != nil {
log.Error(err)
return err
}
// tell the sender the hash so they can quit
c.WriteMessage(websocket.BinaryMessage, append([]byte("hash:"), hash256...))
case 3:
// receive the hash from the sender so we can check it and quit
log.Debugf("got hash: %x", message)
if bytes.Equal(hash256, message) {
c.WriteMessage(websocket.BinaryMessage, []byte("ok"))
// open directory
if fstats.IsDir {
err = zipper.UnzipFile(fstats.SentName, ".")
if DebugLevel != "debug" {
os.Remove(fstats.SentName)
}
} else {
err = nil
}
if err == nil {
transferRate := float64(fstats.Size) / 1000000.0 / transferTime.Seconds()
transferType := "MB/s"
if transferRate < 1 {
transferRate = float64(fstats.Size) / 1000.0 / transferTime.Seconds()
transferType = "kB/s"
}
folderOrFile := "file"
if fstats.IsDir {
folderOrFile = "folder"
}
fmt.Fprintf(os.Stderr, "\nReceived %s written to %s (%2.1f %s)\n", folderOrFile, fstats.Name, transferRate, transferType)
}
return err
} else {
c.WriteMessage(websocket.BinaryMessage, []byte("not"))
return errors.New("file corrupted")
}
default: default:
return fmt.Errorf("unknown step") return fmt.Errorf("unknown step")
} }

View File

@ -19,12 +19,12 @@ const (
pingPeriod = (pongWait * 9) / 10 pingPeriod = (pongWait * 9) / 10
// Maximum message size allowed from peer. // Maximum message size allowed from peer.
maxMessageSize = 1024 * 1024 * 16 maxMessageSize = 1024 * 1024 * 1
) )
var upgrader = websocket.Upgrader{ var upgrader = websocket.Upgrader{
ReadBufferSize: 1024 * 1024 * 8, ReadBufferSize: 1024 * 1024 * 32,
WriteBufferSize: 1024 * 1024 * 8, WriteBufferSize: 1024 * 1024 * 32,
} }
// connection is an middleman between the websocket connection and the hub. // connection is an middleman between the websocket connection and the hub.

View File

@ -44,6 +44,7 @@ func Send(done chan struct{}, c *websocket.Conn, fname string, codephrase string
func send(c *websocket.Conn, fname string, codephrase string) (err error) { func send(c *websocket.Conn, fname string, codephrase string) (err error) {
var f *os.File var f *os.File
var fstats models.FileStats var fstats models.FileStats
var fileHash []byte
// normalize the file name // normalize the file name
fname, err = filepath.Abs(fname) fname, err = filepath.Abs(fname)
@ -156,7 +157,7 @@ func send(c *websocket.Conn, fname string, codephrase string) (err error) {
fmt.Fprintf(os.Stderr, "Sending...\n") fmt.Fprintf(os.Stderr, "Sending...\n")
// send file, compure hash simultaneously // send file, compure hash simultaneously
buffer := make([]byte, 1024*1024*4) buffer := make([]byte, 1024*512)
bar := progressbar.NewOptions( bar := progressbar.NewOptions(
int(fstats.Size), int(fstats.Size),
progressbar.OptionSetRenderBlankState(true), progressbar.OptionSetRenderBlankState(true),
@ -184,12 +185,12 @@ func send(c *websocket.Conn, fname string, codephrase string) (err error) {
err = errors.Wrap(err, "problem writing message") err = errors.Wrap(err, "problem writing message")
return err return err
} }
// wait for ok // // wait for ok
c.ReadMessage() // c.ReadMessage()
} }
if err != nil { if err != nil {
if err != io.EOF { if err != io.EOF {
fmt.Println(err) log.Error(err)
} }
break break
} }
@ -197,14 +198,18 @@ func send(c *websocket.Conn, fname string, codephrase string) (err error) {
bar.Finish() bar.Finish()
log.Debug("send hash to finish file") log.Debug("send hash to finish file")
fileHash, err := utils.HashFile(fname) fileHash, err = utils.HashFile(fname)
if err != nil { if err != nil {
return err return err
} }
c.WriteMessage(websocket.TextMessage, fileHash)
case 4: case 4:
if !bytes.HasPrefix(message, []byte("hash:")) {
continue
}
c.WriteMessage(websocket.BinaryMessage, fileHash)
message = bytes.TrimPrefix(message, []byte("hash:"))
log.Debugf("[%d] determing whether it went ok", step) log.Debugf("[%d] determing whether it went ok", step)
if bytes.Equal(message, []byte("ok")) { if bytes.Equal(message, fileHash) {
log.Debug("file transfered successfully") log.Debug("file transfered successfully")
fmt.Fprintf(os.Stderr, "\nTransfer complete") fmt.Fprintf(os.Stderr, "\nTransfer complete")
return nil return nil