Add file upload support

config.toml.sample

@@ -61,3 +61,12 @@ prefix_session = "NIL:SESS:ROOM:%s"
 # FileSystem store config.
 # [store]
 # path = "db.json"
+
+# File upload configuration.
+[upload]
+max-memory="32MB"
+max-upload-size="2MB"
+max-age="1year"
+rate-limit-count="10"
+rate-limit-period="1minute"
+rate-limit-burst="1"

go.mod

@@ -3,15 +3,18 @@ module github.com/knadh/niltalk
 go 1.13
 
 require (
+	github.com/alecthomas/units v0.0.0-20190924025748-f65c72e2690d
 	github.com/clementauger/tor-prebuilt v0.0.0-20200815153310-0d7058794224
 	github.com/cretz/bine v0.1.0
 	github.com/go-chi/chi v4.1.0+incompatible
 	github.com/gomodule/redigo v2.0.0+incompatible
 	github.com/gorilla/websocket v1.4.2
+	github.com/karrick/tparse/v2 v2.8.2
 	github.com/knadh/koanf v0.9.1
 	github.com/knadh/stuffbin v1.1.0
 	github.com/kr/pretty v0.1.0 // indirect
 	github.com/spf13/pflag v1.0.5
 	golang.org/x/crypto v0.0.0-20200403201458-baeed622b8d8
+	golang.org/x/time v0.0.0-20200630173020-3af7569d3a1e
 	gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 // indirect
 )

go.sum

@@ -1,5 +1,7 @@
 github.com/BurntSushi/toml v0.3.1 h1:WXkYYl6Yr3qBf1K79EBnL4mak0OimBfB0XUf9Vl28OQ=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
+github.com/alecthomas/units v0.0.0-20190924025748-f65c72e2690d h1:UQZhZ2O0vMHr2cI+DC1Mbh0TJxzA3RcLoMsFw+aXw7E=
+github.com/alecthomas/units v0.0.0-20190924025748-f65c72e2690d/go.mod h1:rBZYJk541a8SKzHPHnH3zbiI+7dagKZ0cgpgrD7Fyho=
 github.com/clementauger/tor-prebuilt v0.0.0-20200815153310-0d7058794224 h1:Bm1RJ6O3xTpatREKCjtK3kCmG8SYDgHNGp/qUy0fYek=
 github.com/clementauger/tor-prebuilt v0.0.0-20200815153310-0d7058794224/go.mod h1:QVD8AVR2PuMTcxIbUiUBgexRjeVUtymMMVNrGagoVA4=
 github.com/cretz/bine v0.1.0 h1:1/fvhLE+fk0bPzjdO5Ci+0ComYxEMuB1JhM4X5skT3g=
@@ -18,6 +20,8 @@ github.com/gorilla/websocket v1.4.2 h1:+/TMaTYc4QFitKJxsQ7Yye35DkWvkdLcvGKqM+x0U
 github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/hashicorp/hcl v1.0.0 h1:0Anlzjpi4vEasTeNFn2mLJgTSwt0+6sfsiTG8qcWGx4=
 github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ=
+github.com/karrick/tparse/v2 v2.8.2 h1:NhvrrB7nXYa0VLn0JKn9L3oG/GZN+LB/+g5QfWE30rU=
+github.com/karrick/tparse/v2 v2.8.2/go.mod h1:OzmKMqNal7LYYHaO/Ie1f/wXmLWAaGKwJmxUFNQCVxg=
 github.com/knadh/koanf v0.9.1 h1:qfcwiF9/Z8buTJ0QXaZvOxJ6eKJmOiiWKP/PktiW5RE=
 github.com/knadh/koanf v0.9.1/go.mod h1:31bzRSM7vS5Vm9LNLo7B2Re1zhLOZT6EQKeodixBikE=
 github.com/knadh/stuffbin v1.1.0 h1:f5S5BHzZALjuJEgTIOMC9NidEnBJM7Ze6Lu1GHR/lwU=
@@ -39,6 +43,8 @@ github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.3.0 h1:TivCn/peBQ7UY8ooIcPgZFpTNSz0Q2U6UrFlUfqbe0Q=
 github.com/stretchr/testify v1.3.0/go.mod h1:M5WIy9Dh21IEIfnGCwXGc5bZfKNJtfHm1UVUgZn+9EI=
+github.com/stretchr/testify v1.4.0 h1:2E4SXV/wtOkTonXsotYi4li6zVWxYlZuYNCXe9XRJyk=
+github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20200403201458-baeed622b8d8 h1:fpnn/HnJONpIu6hkXi1u/7rR0NzilgWr4T0JmWkEitk=
 golang.org/x/crypto v0.0.0-20200403201458-baeed622b8d8/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
@@ -51,8 +57,11 @@ golang.org/x/sys v0.0.0-20200331124033-c3d80250170d h1:nc5K6ox/4lTFbMVSL9WRR81ix
 golang.org/x/sys v0.0.0-20200331124033-c3d80250170d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/text v0.3.0 h1:g61tztE5qeGQ89tm6NTjjM9VPIm088od1l6aSorWRWg=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
+golang.org/x/time v0.0.0-20200630173020-3af7569d3a1e h1:EHBhcS0mlXEAVwNyO2dLfjToGsyY4j24pTs2ScHnX7s=
+golang.org/x/time v0.0.0-20200630173020-3af7569d3a1e/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127 h1:qIbj1fsPNlZgppZ+VLlY7N33q108Sa+fhmuc+sWQYwY=
 gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
 gopkg.in/yaml.v2 v2.2.8 h1:obN1ZagJSUGI0Ek/LBmuj4SNLPfIny3KsKFopxRdj10=
 gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=

handlers.go

@@ -4,13 +4,20 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
+	"fmt"
 	"io/ioutil"
+	"log"
 	"net/http"
+	"strings"
+	"sync"
+	"time"
 
 	"github.com/go-chi/chi"
 	"github.com/gorilla/websocket"
 	"github.com/knadh/niltalk/internal/hub"
+	"github.com/knadh/niltalk/internal/upload"
 	"golang.org/x/crypto/bcrypt"
+	"golang.org/x/time/rate"
 )
 
 const (
@@ -337,3 +344,104 @@ func readJSONReq(r *http.Request, o interface{}) error {
 	}
 	return json.Unmarshal(b, o)
 }
+
+// handleUpload handles file uploads.
+func handleUpload(store *upload.Store) func(w http.ResponseWriter, r *http.Request) {
+
+	type roomLimiter struct {
+		limiter *rate.Limiter
+		expire  time.Time
+	}
+	var mu sync.Mutex
+	roomLimiters := map[string]roomLimiter{}
+	go func() {
+		t := time.NewTicker(store.RlPeriod + (time.Minute))
+		defer t.Stop()
+		for range t.C {
+			now := time.Now()
+			mu.Lock()
+			for k, r := range roomLimiters {
+				if r.expire.Before(now) {
+					delete(roomLimiters, k)
+				}
+			}
+			mu.Unlock()
+		}
+	}()
+
+	return func(w http.ResponseWriter, r *http.Request) {
+		r.ParseMultipartForm(store.MaxUploadSize)
+
+		roomID := chi.URLParam(r, "roomID")
+		mu.Lock()
+		// no defer here becasue file upload can be slow, thus lock for too long
+		x, ok := roomLimiters[roomID]
+		if !ok {
+			x = roomLimiter{
+				limiter: rate.NewLimiter(rate.Every(store.RlPeriod/time.Duration(store.RlCount)), store.RlBurst),
+				expire:  time.Now().Add(time.Minute * 10),
+			}
+			roomLimiters[roomID] = x
+		}
+		x.expire = time.Now().Add(time.Minute * 10)
+		roomLimiters[roomID] = x
+		mu.Unlock()
+		if !x.limiter.Allow() {
+			err := errors.New(http.StatusText(http.StatusTooManyRequests))
+			respondJSON(w, nil, err, http.StatusTooManyRequests)
+			return
+		}
+
+		var ids []string
+		for i := 0; i < 20; i++ {
+			key := fmt.Sprintf("file%v", i)
+			file, handler, err := r.FormFile(key)
+			if err == http.ErrMissingFile {
+				break
+			}
+			if err != nil {
+				continue
+			}
+			defer file.Close()
+			b, err := ioutil.ReadAll(file)
+			if err != nil {
+				continue
+			}
+			mimeType := http.DetectContentType(b)
+			if mimeType == "image/gif" || mimeType == "image/jpeg" || mimeType == "image/png" {
+				name := handler.Filename
+				up, err := store.Add(name, mimeType, b)
+				if err != nil {
+					continue
+				}
+				ids = append(ids, fmt.Sprintf("%v_%v", up.ID, up.Name))
+			}
+		}
+
+		respondJSON(w, struct {
+			IDs []string `json:"ids"`
+		}{ids}, nil, http.StatusOK)
+	}
+}
+
+// handleUploaded uploaded files display.
+func handleUploaded(store *upload.Store) func(w http.ResponseWriter, r *http.Request) {
+	maxAgeHeader := fmt.Sprintf("max-age=%v", int64(store.MaxAge/time.Second))
+	return func(w http.ResponseWriter, r *http.Request) {
+		fileID := chi.URLParam(r, "fileID")
+		fileID = strings.Split(fileID, "_")[0]
+		up, err := store.Get(fileID)
+		if err != nil {
+			log.Println(err)
+			respondJSON(w, nil, err, http.StatusNotFound)
+			return
+		}
+		w.Header().Add("Content-Type", up.MimeType)
+		w.Header().Add("Content-Length", fmt.Sprint(len(up.Data)))
+		if store.MaxAge > 0 {
+			w.Header().Add("Cache-Control", maxAgeHeader)
+		}
+		w.WriteHeader(http.StatusOK)
+		w.Write(up.Data)
+	}
+}

internal/hub/hub.go

@@ -14,6 +14,7 @@ import (
 const (
 	TypeTyping          = "typing"
 	TypeMessage         = "message"
+	TypeUpload          = "upload"
 	TypePeerList        = "peer.list"
 	TypePeerInfo        = "peer.info"
 	TypePeerJoin        = "peer.join"

internal/hub/peer.go

@@ -127,7 +127,36 @@ func (p *Peer) processMessage(b []byte) {
 			// TODO: Respond
 			return
 		}
-		p.room.Broadcast(p.room.makeMessagePayload(msg, p), true)
+		p.room.Broadcast(p.room.makeMessagePayload(msg, p, m.Type), true)
+
+	case TypeUpload:
+		// Check rate limits and update counters.
+		now := time.Now()
+		if p.numMessages > 0 {
+			if (p.numMessages%p.room.hub.cfg.RateLimitMessages+1) >= p.room.hub.cfg.RateLimitMessages &&
+				time.Since(p.lastMessage) < p.room.hub.cfg.RateLimitInterval {
+				p.room.hub.Store.RemoveSession(p.ID, p.room.ID)
+				p.writeWSControl(websocket.CloseMessage,
+					websocket.FormatCloseMessage(websocket.CloseNormalClosure, TypePeerRateLimited))
+				p.ws.Close()
+				return
+			}
+		}
+		p.lastMessage = now
+		p.numMessages++
+
+		msgs, ok := m.Data.([]interface{})
+		if !ok {
+			// TODO: Respond
+			return
+		}
+		for _, msg := range msgs {
+			x, ok := msg.(string)
+			if !ok {
+				continue
+			}
+			p.room.Broadcast(p.room.makeMessagePayload(x, p, m.Type), true)
+		}
 
 	// "Typing" status.
 	case TypeTyping:

internal/hub/room.go

@@ -258,13 +258,13 @@ func (r *Room) makePeerUpdatePayload(p *Peer, peerUpdateType string) []byte {
 }
 
 // makeMessagePayload prepares a chat message.
-func (r *Room) makeMessagePayload(msg string, p *Peer) []byte {
+func (r *Room) makeMessagePayload(msg string, p *Peer, typ string) []byte {
 	d := payloadMsgChat{
 		PeerID:     p.ID,
 		PeerHandle: p.Handle,
 		Msg:        msg,
 	}
-	return r.makePayload(d, TypeMessage)
+	return r.makePayload(d, typ)
 }
 
 // makePayload prepares a message payload.