Skip to content

Instantly share code, notes, and snippets.

@tsilvers

tsilvers/upload.go

Last active August 1, 2024 13:57
Show Gist options
  • Save tsilvers/5f827fb11aee027e22c6b3102ebcc497 to your computer and use it in GitHub Desktop.
Save tsilvers/5f827fb11aee027e22c6b3102ebcc497 to your computer and use it in GitHub Desktop.
Download ZIP
Upload files to a Go server using web sockets.
Raw
upload.go
package main
import (
"encoding/json"
"fmt"
"io/ioutil"
"log"
"net/http"
"os"
"time"
"github.com/gorilla/websocket"
)
const HandshakeTimeoutSecs = 10
type UploadHeader struct {
Filename string
Size int
}
type UploadStatus struct {
Code int `json:"code,omitempty"`
Status string `json:"status,omitempty"`
Pct *int `json:"pct,omitempty"` // File processing AFTER upload is done.
pct int
}
type wsConn struct {
conn *websocket.Conn
}
func upload(w http.ResponseWriter, r *http.Request) {
wsc := wsConn{}
var err error
// Open websocket connection.
upgrader := websocket.Upgrader{HandshakeTimeout: time.Second * HandshakeTimeoutSecs}
wsc.conn, err = upgrader.Upgrade(w, r, nil)
if err != nil {
fmt.Println("Error on open of websocket connection:", err)
return
}
defer wsc.conn.Close()
// Get upload file name and length.
header := new(UploadHeader)
mt, message, err := wsc.conn.ReadMessage()
if err != nil {
fmt.Println("Error receiving websocket message:", err)
return
}
if mt != websocket.TextMessage {
wsc.sendStatus(400, "Invalid message received, expecting file name and length")
return
}
if err := json.Unmarshal(message, header); err != nil {
wsc.sendStatus(400, "Error receiving file name and length: "+err.Error())
return
}
if len(header.Filename) == 0 {
wsc.sendStatus(400, "Filename cannot be empty")
return
}
if header.Size == 0 {
wsc.sendStatus(400, "Upload file is empty")
return
}
// Create temp file to save file.
var tempFile *os.File
if tempFile, err = ioutil.TempFile("", "websocket_upload_"); err != nil {
wsc.sendStatus(400, "Could not create temp file: "+err.Error())
return
}
defer func() {
tempFile.Close()
// *** IN PRODUCTION FILE SHOULD BE REMOVED AFTER PROCESSING ***
// _ = os.Remove(tempFile.Name())
}()
// Read file blocks until all bytes are received.
bytesRead := 0
for {
mt, message, err := wsc.conn.ReadMessage()
if err != nil {
wsc.sendStatus(400, "Error receiving file block: "+err.Error())
return
}
if mt != websocket.BinaryMessage {
if mt == websocket.TextMessage {
if string(message) == "CANCEL" {
wsc.sendStatus(400, "Upload canceled")
return
}
}
wsc.sendStatus(400, "Invalid file block received")
return
}
tempFile.Write(message)
bytesRead += len(message)
if bytesRead == header.Size {
tempFile.Close()
break
}
wsc.requestNextBlock()
}
// *****************************
// *** Process uploaded file ***
// *****************************
for i := 0; i <= 10; i++ {
wsc.sendPct(i * 10)
time.Sleep(time.Second * 1)
}
wsc.sendStatus(200, "File upload successful: "+fmt.Sprintf("%s (%d bytes)", tempFile.Name(), bytesRead))
}
func (wsc wsConn) requestNextBlock() {
wsc.conn.WriteMessage(websocket.TextMessage, []byte("NEXT"))
}
func (wsc wsConn) sendStatus(code int, status string) {
if msg, err := json.Marshal(UploadStatus{Code: code, Status: status}); err == nil {
wsc.conn.WriteMessage(websocket.TextMessage, msg)
}
}
func (wsc wsConn) sendPct(pct int) {
stat := UploadStatus{pct: pct}
stat.Pct = &stat.pct
if msg, err := json.Marshal(stat); err == nil {
wsc.conn.WriteMessage(websocket.TextMessage, msg)
}
}
func showJS(w http.ResponseWriter, r *http.Request) {
fmt.Fprintf(w, "%s", webPage)
}
func main() {
http.HandleFunc("/", showJS)
http.HandleFunc("/upload", upload)
log.Fatal(http.ListenAndServe(":80", nil))
}
var webPage = `
<!DOCTYPE html>
<html>
<head>
<meta charset="utf-8">
</head>
<body>
<br>
<input type="file" id="uploadFile" name="file" />
<br><br>
<span id="uploadButton">
<button id="upload">Upload</button>
</span>
<span id="cancelButton" style="display: none;">
<button id="cancel" style="color: red">CANCEL</button>
</span>
<br>
<br>
Upload %: <output id="pct"></output>
<br>
<br>
Status Messages: <output id="list"></output>
</body>
<script>
var blockSize = 1024 * 1024;
var ws;
var filePos;
var file;
var reader;
var blob;
var cancel = false;
function readBlob() {
var first = filePos;
var last = first + blockSize
if (last > file.size) {
last == file.size;
}
blob = file.slice(first, last);
reader.readAsArrayBuffer(blob);
}
document.querySelector('#upload').addEventListener('click', function(evt) {
document.getElementById('list').innerHTML = '';
document.getElementById('pct').innerHTML = '';
filePos = 0;
reader = new FileReader();
cancel = false;
if (evt.target.tagName.toLowerCase() == 'button') {
var files = document.getElementById('uploadFile').files;
if (!files.length) {
alert('Please select a file!');
return;
}
file = files[0];
document.getElementById('uploadButton').style.display = 'none';
document.getElementById('cancelButton').style.display = 'block';
// Open and configure websocket connection.
ws = new WebSocket("ws://localhost/upload");
ws.binaryType = 'arraybuffer';
// Send filename and size to the server when the connection is opened.
ws.onopen = function(evt) {
header = '{"filename":"' + file.name + '","size":' + file.size + '}';
ws.send(header);
};
// Initiate the file transfer by reading the first block from disk.
readBlob();
// Send the next file block to the server once it's read from disk.
reader.onloadend = function(evt) {
if (evt.target.readyState == FileReader.DONE) {
ws.send(blob);
filePos += blob.size;
document.getElementById('pct').innerHTML = filePos / file.size * 100.0;
if (filePos >= file.size || cancel) {
document.getElementById('cancelButton').style.display = 'none';
document.getElementById('uploadButton').style.display = 'block';
}
}
};
// Process message sent from server.
ws.onmessage = function(e) {
// Server only sends text messages.
if (typeof e.data === "string") {
// "NEXT" message confirms the server received the last block.
if (e.data === "NEXT") {
// If we're not cancelling the upload, read the next file block from disk.
if (cancel) {
document.getElementById('cancelButton').style.display = 'none';
document.getElementById('uploadButton').style.display = 'block';
} else {
readBlob();
}
// Otherwise, message is a status update (json).
} else {
document.getElementById('list').innerHTML = '<ul>' + e.data + '</ul>';
}
}
};
ws.onclose = function(evt) {
ws = null;
};
ws.onerror = function(evt) {
ws.close();
ws = null;
return false;
};
}
}, false);
document.querySelector('#cancel').addEventListener('click', function(evt) {
cancel = true;
ws.send("CANCEL");
}, false);
</script>
</html>
`
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment