ooni-probe-cli/internal/engine/experiment/ndt7/download.go
Simone Basso d57c78bc71
chore: merge probe-engine into probe-cli (#201)
This is how I did it:

1. `git clone https://github.com/ooni/probe-engine internal/engine`

2. ```
(cd internal/engine && git describe --tags)
v0.23.0
```

3. `nvim go.mod` (merging `go.mod` with `internal/engine/go.mod`

4. `rm -rf internal/.git internal/engine/go.{mod,sum}`

5. `git add internal/engine`

6. `find . -type f -name \*.go -exec sed -i 's@/ooni/probe-engine@/ooni/probe-cli/v3/internal/engine@g' {} \;`

7. `go build ./...` (passes)

8. `go test -race ./...` (temporary failure on RiseupVPN)

9. `go mod tidy`

10. this commit message

Once this piece of work is done, we can build a new version of `ooniprobe` that
is using `internal/engine` directly. We need to do more work to ensure all the
other functionality in `probe-engine` (e.g. making mobile packages) are still WAI.

Part of https://github.com/ooni/probe/issues/1335
2021-02-02 12:05:47 +01:00

74 lines
1.5 KiB
Go

package ndt7
import (
"context"
"io"
"io/ioutil"
"time"
"github.com/gorilla/websocket"
)
type downloadManager struct {
conn mockableConn
maxMessageSize int64
maxRuntime time.Duration
measureInterval time.Duration
onJSON callbackJSON
onPerformance callbackPerformance
}
func newDownloadManager(
conn mockableConn, onPerformance callbackPerformance,
onJSON callbackJSON,
) downloadManager {
return downloadManager{
conn: conn,
maxMessageSize: paramMaxMessageSize,
maxRuntime: paramMaxRuntime,
measureInterval: paramMeasureInterval,
onJSON: onJSON,
onPerformance: onPerformance,
}
}
func (mgr downloadManager) run(ctx context.Context) error {
var total int64
start := time.Now()
if err := mgr.conn.SetReadDeadline(start.Add(mgr.maxRuntime)); err != nil {
return err
}
mgr.conn.SetReadLimit(mgr.maxMessageSize)
ticker := time.NewTicker(mgr.measureInterval)
defer ticker.Stop()
for ctx.Err() == nil {
kind, reader, err := mgr.conn.NextReader()
if err != nil {
return err
}
if kind == websocket.TextMessage {
data, err := ioutil.ReadAll(reader)
if err != nil {
return err
}
total += int64(len(data))
if err := mgr.onJSON(data); err != nil {
return err
}
continue
}
n, err := io.Copy(ioutil.Discard, reader)
if err != nil {
return err
}
total += int64(n)
select {
case now := <-ticker.C:
mgr.onPerformance(now.Sub(start), total)
default:
// NOTHING
}
}
return nil
}