github.com/LQR471814/websocket-ftp/server@v0.4.0/server.go (about) 1 package server 2 3 import ( 4 "bufio" 5 "bytes" 6 "encoding/json" 7 "io" 8 "log" 9 "net" 10 "net/http" 11 "strings" 12 13 "github.com/gorilla/websocket" 14 ) 15 16 func sendJSON(c *websocket.Conn, obj interface{}) { 17 msg, err := json.Marshal(obj) 18 if err != nil { 19 panic(err) 20 } 21 22 c.WriteMessage(websocket.TextMessage, msg) 23 } 24 25 func outputWriter(s *WSFTPServer, f File, t *Transfer) { 26 _, ok := t.Output[f.ID()] 27 if !ok { 28 t.Output[f.ID()] = bytes.NewBuffer(nil) 29 } 30 31 var writtenBytes int64 = 0 32 w := bufio.NewWriterSize(t.Output[f.ID()], 1024*1024*50) //? Buffsize = 50 mB 33 34 for data := range t.dataChan { 35 w.Write(data) 36 writtenBytes += int64(len(data)) 37 if writtenBytes >= f.Size { 38 if s.Config.Verbose { 39 log.Printf("--> DONE: Wrote %v to output\n", f.Name) 40 } 41 w.Flush() 42 return 43 } 44 } 45 } 46 47 func Handler(s *WSFTPServer, w http.ResponseWriter, r *http.Request) { 48 var updateRatio = 24 49 if s.Config.Verbose { 50 updateRatio = 4 51 } 52 53 conn, err := upgrader.Upgrade(w, r, nil) //? Event: onpeerconnect 54 if err != nil { 55 panic(err) 56 } 57 58 transfer := &Transfer{ 59 Data: TransferMetadata{ 60 From: conn.RemoteAddr().(*net.TCPAddr).IP, 61 }, 62 ID: s.ids.Fetch(), 63 State: TransferState{Number: INITIAL}, 64 Output: make(map[string]io.Writer), 65 conn: conn, 66 } 67 68 s.Transfers[transfer.ID] = transfer 69 eventHandler(s, transfer, peerConnect) 70 71 var updateNext int64 = 0 72 73 for { 74 msgType, contents, err := conn.ReadMessage() 75 if err != nil { 76 if strings.Contains(err.Error(), "close") { 77 return 78 } 79 panic(err) 80 } 81 82 switch msgType { 83 case websocket.TextMessage: 84 reqs := &FileRequests{} 85 json.Unmarshal(contents, reqs) 86 transfer.Data.Files = reqs.Files 87 eventHandler(s, transfer, recvRequests) 88 case websocket.BinaryMessage: 89 f := transfer.Data.Files[transfer.State.CurrentFile] 90 91 updateOffset := f.Size / int64(updateRatio) 92 93 transfer.dataChan <- contents 94 transfer.State.Received += int64(len(contents)) 95 96 if transfer.State.Received >= updateNext { 97 if s.Config.Handlers != nil { 98 s.Config.Handlers.OnTransferUpdate(transfer) 99 } 100 updateNext += updateOffset 101 } 102 103 if transfer.State.Received >= f.Size { 104 eventHandler(s, transfer, recvDone) 105 updateNext = 0 106 } 107 } 108 } 109 } 110 111 func eventHandler(s *WSFTPServer, t *Transfer, event Event) { 112 cell, ok := EventStateMatrix[event][t.State.Number] 113 if s.Config.Verbose { 114 log.Println("Event", event, "State", t.State, cell) 115 } 116 117 if !ok { 118 panic("Invalid FileTransfer state") 119 } 120 121 t.State.Number = cell.NewState 122 for _, action := range cell.Actions { 123 actionHandler(s, t, action) 124 } 125 } 126 127 func actionHandler(s *WSFTPServer, t *Transfer, action Action) { 128 switch action { 129 case DisplayFileRequests: 130 var accept = true 131 if s.Config.Handlers != nil { 132 accept = <-s.Config.Handlers.OnTransferRequest(t) 133 } 134 if accept { 135 eventHandler(s, t, userAccept) 136 return 137 } 138 eventHandler(s, t, userDeny) 139 case IncrementFileIndex: 140 s.Config.Handlers.OnTransferComplete(t, t.Data.Files[t.State.CurrentFile]) 141 t.State.CurrentFile += 1 142 t.State.Received = 0 143 case SendStartSignal: 144 sendJSON(t.conn, Signal{Type: "start"}) 145 case SendExitSignal: 146 sendJSON(t.conn, Signal{Type: "exit"}) 147 case SendFinishedSignal: 148 sendJSON(t.conn, Signal{Type: "complete"}) 149 case StartFileWriter: 150 t.dataChan = make(chan []byte) 151 f := t.Data.Files[t.State.CurrentFile] 152 go outputWriter(s, f, t) 153 case StopFileWriter: 154 close(t.dataChan) 155 case RecvDoneHandler: 156 if t.State.CurrentFile >= len(t.Data.Files) { 157 s.Config.Handlers.OnAllTransfersComplete(t) 158 return 159 } 160 actionHandler(s, t, StartFileWriter) 161 actionHandler(s, t, SendStartSignal) 162 } 163 }