CODEWITHSUNDEEP
performancedownloadgothrottlingrate
Royalty
Royalty

Reputation: 45

How to limit download speed with Go?

I'm currently developing a download server in Go. I need to limit the download speed of users to 100KB/s.

This was my code:

func serveFile(w http.ResponseWriter, r *http.Request) {
    fileID := r.URL.Query().Get("fileID")
    if len(fileID) != 0 {
        w.Header().Set("Content-Disposition", "attachment; filename=filename.txt")
        w.Header().Set("Content-Type", r.Header.Get("Content-Type"))
        w.Header().Set("Content-Length", r.Header.Get("Content-Length"))

        file, err := os.Open(fmt.Sprintf("../../bin/files/test.txt"))
        defer file.Close()
        if err != nil {
            http.NotFound(w, r)
            return
        }
        io.Copy(w, file)
    } else {
        io.WriteString(w, "Invalid request.")
    }
}

Then I found a package on github and my code became the following:

func serveFile(w http.ResponseWriter, r *http.Request) {
    fileID := r.URL.Query().Get("fileID")
    if len(fileID) != 0 {
        w.Header().Set("Content-Disposition", "attachment; filename=Wiki.png")
        w.Header().Set("Content-Type", r.Header.Get("Content-Type"))
        w.Header().Set("Content-Length", r.Header.Get("Content-Length"))

        file, err := os.Open(fmt.Sprintf("../../bin/files/test.txt"))
        defer file.Close()
        if err != nil {
            http.NotFound(w, r)
            return
        }
        bucket := ratelimit.NewBucketWithRate(100*1024, 100*1024)
        reader := bufio.NewReader(file)
        io.Copy(w, ratelimit.Reader(reader, bucket))
    } else {
        io.WriteString(w, "Invalid request.")
    }
}

But I'm getting this error:

Corrupted Content Error

The page you are trying to view cannot be shown because an error in the data transmission was detected.

Here's my code on the Go playground: http://play.golang.org/p/ulgXQl4eQO

Upvotes: 3

Views: 3228

Answers (2)

Dave C
Dave C

Reputation: 7878

Rather than mucking around with getting the correct the content type and length headers yourself it'd probably be much better to use http.ServeContent which will do that for you (as well as support "If-Modified-Since", range requests, etc. If you can supply an "ETag" header it can also handle "If-Range" and "If-None-Match" requests as well).

As mentioned previously, it's often preferable to limit on the write side but it's awkward to wrap an http.ResponseWriter since various http functions also check for optional interfaces such as http.Flusher and http.Hijacker. It's much easier to wrap the io.ReadSeeker that ServeContent needs.

For example, something like this perhaps:

func pathFromID(fileID string) string {
    // replace with whatever logic you need
    return "../../bin/files/test.txt"
}

// or more verbosely you could call this a "limitedReadSeeker"
type lrs struct {
    io.ReadSeeker
    // This reader must not buffer but just do something simple
    // while passing through Read calls to the ReadSeeker
    r io.Reader
}

func (r lrs) Read(p []byte) (int, error) {
    return r.r.Read(p)
}

func newLRS(r io.ReadSeeker, bucket *ratelimit.Bucket) io.ReadSeeker {
    // Here we know/expect that a ratelimit.Reader does nothing
    // to the Read calls other than add delays so it won't break
    // any io.Seeker calls.
    return lrs{r, ratelimit.Reader(r, bucket)}
}

func serveFile(w http.ResponseWriter, req *http.Request) {
    fileID := req.URL.Query().Get("fileID")
    if len(fileID) == 0 {
        http.Error(w, "invalid request", http.StatusBadRequest)
        return
    }

    path := pathFromID(fileID)
    file, err := os.Open(path)
    if err != nil {
        http.NotFound(w, req)
        return
    }
    defer file.Close()
    fi, err := file.Stat()
    if err != nil {
        http.Error(w, "blah", 500) // XXX fixme
        return
    }

    const (
        rate     = 100 << 10
        capacity = 100 << 10
    )
    // Normally we'd prefer to limit the writer but it's awkward to wrap
    // an http.ResponseWriter since it may optionally also implement
    // http.Flusher, or http.Hijacker.
    bucket := ratelimit.NewBucketWithRate(rate, capacity)
    lr := newLRS(file, bucket)
    http.ServeContent(w, req, path, fi.ModTime(), lr)
}

Upvotes: 2

Caleb
Caleb

Reputation: 9458

I'm not seeing the error, but I did notice some issues with the code. For this:

w.Header().Set("Content-Type", r.Header.Get("Content-Type"))

You should use the mime package's:

func TypeByExtension(ext string) string

To determine the content type. (if you end up with the empty string default to application/octet-stream)

For:

w.Header().Set("Content-Length", r.Header.Get("Content-Length"))

You need to get the content length from the file itself. By using the request content length, for a GET this basically ends up as a no-op, but for a POST you're sending back the wrong length, which might explain the error you're seeing. After you open the file, do this:

fi, err := file.Stat()
if err != nil {
    http.Error(w, err.Error(), 500)
    return
}
w.Header().Set("Content-Length", fmt.Sprint(fi.Size()))

One final thing, when you open the file, if there's an error, you don't need to close the file handle. Do it like this instead:

file, err := os.Open(...)
if err != nil {
    http.NotFound(w, r)
    return
}
defer file.Close()

Upvotes: 1

Related Questions