Merge bba644c into db6ef3f

m-lab · Feb 5, 2020 · 38c161b · 38c161b
2 parents db6ef3f + bba644c
commit 38c161b
Show file tree

Hide file tree

Showing 3 changed files with 32 additions and 15 deletions.
diff --git a/internal/download/download.go b/internal/download/download.go
@@ -4,6 +4,8 @@ package download
 import (
 	"context"
 	"encoding/json"
+	"io"
+	"io/ioutil"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -34,11 +36,31 @@ func Run(ctx context.Context, conn websocketx.Conn, ch chan<- spec.Measurement)
 		if err != nil {
 			return err
 		}
-		mtype, mdata, err := conn.ReadMessage()
+		mtype, reader, err := conn.NextReader()
 		if err != nil {
 			return err
 		}
-		total += int64(len(mdata))
+		if mtype == websocket.TextMessage {
+			mdata, err := ioutil.ReadAll(reader)
+			if err != nil {
+				return err
+			}
+			total += int64(len(mdata))
+			var measurement spec.Measurement
+			err = json.Unmarshal(mdata, &measurement)
+			if err != nil {
+				return err
+			}
+			measurement.Origin = spec.OriginServer
+			measurement.Test = spec.TestDownload
+			ch <- measurement
+		} else {
+			count, err := io.Copy(ioutil.Discard, reader)
+			if err != nil {
+				return err
+			}
+			total += int64(count)
+		}
 		now := time.Now()
 		if now.Sub(prev) > params.UpdateInterval {
 			prev = now
@@ -53,17 +75,6 @@ func Run(ctx context.Context, conn websocketx.Conn, ch chan<- spec.Measurement)
 			}
 			// FALLTHROUGH
 		}
-		if mtype != websocket.TextMessage {
-			continue
-		}
-		var measurement spec.Measurement
-		err = json.Unmarshal(mdata, &measurement)
-		if err != nil {
-			return err
-		}
-		measurement.Origin = spec.OriginServer
-		measurement.Test = spec.TestDownload
-		ch <- measurement
 	}
 	return nil // this is how success looks like
 }
diff --git a/internal/upload/upload.go b/internal/upload/upload.go
@@ -4,6 +4,8 @@ package upload
 import (
 	"context"
 	"errors"
+	"io"
+	"io/ioutil"
 	"math/rand"
 	"time"
 
@@ -50,13 +52,16 @@ func ignoreIncoming(conn websocketx.Conn) error {
 		if err != nil {
 			return err
 		}
-		mtype, _, err := conn.ReadMessage()
+		mtype, reader, err := conn.NextReader()
 		if err != nil {
 			return err
 		}
 		if mtype != websocket.TextMessage {
 			return errNonTextMessage
 		}
+		if _, err := io.Copy(ioutil.Discard, reader); err != nil {
+			return err
+		}
 	}
 }
 

diff --git a/internal/websocketx/websocketx.go b/internal/websocketx/websocketx.go
@@ -2,6 +2,7 @@
 package websocketx
 
 import (
+	"io"
 	"time"
 
 	"github.com/gorilla/websocket"
@@ -10,7 +11,7 @@ import (
 // Conn is the interface of a websocket.Conn used for mocking.
 type Conn interface {
 	Close() error
-	ReadMessage() (messageType int, p []byte, err error)
+	NextReader() (messageType int, reader io.Reader, err error)
 	SetReadLimit(limit int64)
 	SetReadDeadline(t time.Time) error
 	SetWriteDeadline(t time.Time) error