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  }