DEV Community: Viktor Nikolaiev

On-the-Fly Content Type Sniffing and Validation in Go

Viktor Nikolaiev — Wed, 12 Feb 2025 14:58:51 +0000

Validating, sanitizing, and safely hosting user-generated content is a large and complex topic. Today we'll look at just a single aspect of it — using magic bytes to sniff the content type of user uploaded files and reject ones that do not match our validation rules.

This post focuses on two things:

How to sniff the content type without buffering an entire file in memory
How to make the code ergonomic and reusable

Using http.DetectContentType

There is a http.DetectContentType function in the standard library that does exactly what we need. Here's its full description from the documentation:

DetectContentType implements the algorithm described at https://mimesniff.spec.whatwg.org/ to determine the Content-Type of the given data. It considers at most the first 512 bytes of data. DetectContentType always returns a valid MIME type: if it cannot determine a more specific one, it returns "application/octet-stream".

Let's see how we could use this in practice. Consider a typical file upload handler that copies files directly to S3:

func handleUpload(w http.ResponseWriter, r *http.Request) {
    // Validate request headers
    // ...

    // Copy request body right into S3
    uploader := s3manager.NewUploader(sess)

    _, err := uploader.Upload(&s3manager.UploadInput{
        Bucket: aws.String("my-bucket"),
        Key:    aws.String("filename.jpg"),
        Body:   r.Body,
    })
    if err != nil {
        w.WriteHeader(500)
    }
}

Let's think about what it would look like if we manually used the http.DetectContentType to only allow uploading images. Remember, we don't want to buffer the entire file in memory:

    // Read the first chunk of the request body
    var first512 [512]byte
    n, err := io.ReadFull(r.Body, first512[:])
    if err != nil && !errors.Is(err, io.ErrUnexpectedEOF) && !errors.Is(err, io.EOF) {
        w.WriteHeader(500)
        return
    }

    // Detect and validate the content type
    contentType := http.DetectContentType(first512[:n])
    if !strings.HasPrefix(contentType, "image/") {
        w.WriteHeader(400)
        return
    }

    // Reassemble the request body back
    reqBody := io.MultiReader(bytes.NewReader(first512[:n]), r.Body)

While this works, it has several drawbacks:

The code is verbose and error-prone
Error handling is scattered and difficult to maintain
We'd need to repeat this pattern in every upload handler

Let's encapsulate this logic into a reusable component that handles the complexity for us.

What Would We Want Instead?

Let's first look at what we'd want to achieve. We'd like to have a custom reader wrapper — NewContentTypeReader that automatically detects the content type as the body is being read and calls a user-provided callback. And user's callback would do all necessary validations and optionally return an error. Our HTTP handler would look like this:

var ErrNotImage = errors.New("not an image")

func handleUpload(w http.ResponseWriter, r *http.Request) {
    // Validate request headers
    // ...

    // A custom reader that detects and validates the content type
    reqBody := NewContentTypeReader(r.Body, func(contentType string) error {
        if !strings.HasPrefix(contentType, "image/") {
            return ErrNotImage
        }
        return nil
    })

    // Copy request body right into S3
    uploader := s3manager.NewUploader(sess)

    _, err := uploader.Upload(&s3manager.UploadInput{
        Bucket: aws.String("my-bucket"),
        Key:    aws.String("filename.jpg"),
        Body:   reqBody,
    })
    if errors.Is(err, ErrNotImage) {
        // Caught the error from the custom reader
        w.WriteHeader(400)
        return
    }
    if err != nil {
        w.WriteHeader(500)
    }
}

In this code, any attempt to read from the reqBody would fail with ErrNotImage if the request body does not look like an image.

Implementing the Reader Wrapper

Now after we've seen how we want to use our wrapper, let's look at how to implement it. We could use the same approach as above, with io.MultiReader, though it would cause some complications in the cases when it's impossible to read the first 512 bytes because of some non-fatal error (like i/o timeout).

Instead, our wrapper just proxies all Read calls right to the original reader, but also accumulates an internal buffer until there's enough data to call the http.DetectContentType function.

type ctReader struct {
    buf     []byte
    reader  io.Reader
    handler func(contentType string) error
}

// NewContentTypeReader returns a reader that sniffs the content type and passes it to the handler.
func NewContentTypeReader(r io.Reader, handler func(contentType string) error) io.Reader {
    return &ctReader{
        reader:  r,
        handler: handler,
    }
}

func (r *ctReader) Read(p []byte) (n int, err error) {
    n, err = r.reader.Read(p)

    if r.handler != nil {
        // Accumulate the buffer
        r.buf = append(r.buf, p[:min(512, n)]...)

        // Buffer is large enough or EOF reached
        if len(r.buf) >= 512 || errors.Is(err, io.EOF) {
            contentType := http.DetectContentType(r.buf)
            if err2 := r.handler(contentType); err2 != nil {
                err = err2 // replace the original error
            }

            // Make sure we don't call the handler again
            r.handler = nil
            r.buf = nil
        }
    }
    return
}

Note: The code, as it's written above, can buffer up to 1023 bytes. While it's possible theoretically (511b read + 512b read), it's very unlikely to happen in practice, where each read operation is usually several Kb long.

Conclusion

We've built a reusable solution for content-type detection during file uploads that:

Validates files on-the-fly without buffering them entirely in memory
Integrates with Go's io.Reader interface and standard libraries
Works naturally with cloud storage services like Amazon S3

While magic byte detection is reliable for most common file types, remember that it's just one layer of defense. For production systems, combine it with other security measures like file size limits, malware scanning, and proper access controls.

Originally published at destel.dev

Real-Time Batching in Go

Viktor Nikolaiev — Mon, 03 Feb 2025 17:15:32 +0000

Batching is a well-known optimization technique. You see it everywhere: batch inserts in databases, MGET/MSET in Redis, various bulk APIs. The benefits are clear — it's faster, cheaper, and less rate-limited. These benefits usually come at the cost of slightly more complex code and some boilerplate.

But how do we cleanly batch something we can't see yet? Something that's arriving in real time.

A Real-World Example

Imagine an application that needs to update users' last_active_at timestamp in the database each time the user interacts with the system. To make this example more extreme, let's assume that the app updates this timestamp on every HTTP request from a user. Our function might look like this:

func UpdateUserTimestamp(ctx context.Context, userID int) error {
    // UPDATE users SET last_active_at = NOW() WHERE id = userID
}

With thousands of concurrent requests, we're hitting the database just as many times per second. Each update is tiny, but together they can create unnecessary database pressure. We'd want to group together the updates happening at the same time. Like this:

UPDATE users SET last_active_at=NOW() WHERE id IN (17, 25,31)

The challenge is that UpdateUserTimestamp only sees one user ID at a time, making batching impossible at this level.

We can solve this by having UpdateUserTimestamp send each ID to a shared queue (implemented as a Go channel). A background worker can then read from this queue, accumulate IDs into batches, and send bulk updates to the database.

Making it Developer-Friendly

While the basic concept is straightforward, the devil is in the details. We want the UpdateUserTimestamp behave exactly as before: accept a single user ID, respect context cancellation, and return appropriate errors. Let's build a solution that handles all these requirements.

Let's first deal with error handling. When UpdateUserTimestamp sends a user ID to the worker, how does it know if that specific update succeeded or failed? Our solution is to create a dedicated ReplyTo channel for each call. The worker uses this channel to report success or failure for that specific update. This naturally provides the blocking behavior we want — each function call waits for its own result via its unique channel. From the developer's perspective, the function appears to simply block while performing a query.

Let's visualize how two concurrent UpdateUserTimestamp calls flow through the system. This sequence diagram shows the key interactions:

Full-sized SVG diagram

Now let's implement the green part of the diagram in the code:

// This type represents a single request to the worker
type updateUserTimestampRequest struct {
    userID  int
    ReplyTo chan error
}

// This is the queue of user IDs to update.
var updateUserTimestampQueue = make(chan updateUserTimestampRequest)

// UpdateUserTimestamp is the public API for updating the last_active_at column in the users table
func UpdateUserTimestamp(ctx context.Context, userID int) error {
    // Prepare a request to the worker.
    // A ReplyTo channel is used by the worker to send us back a result.
    req := updateUserTimestampRequest{
        userID:  userID,
        ReplyTo: make(chan error),
    }

    // Send request to the worker
    updateUserTimestampQueue <- req

    // Block and wait for the result
    err := <-req.ReplyTo
    return err
}

What About the Context Support?

The code above ignores the context passed to the function. Let's fix this. We won't use this context for database query timeouts — the worker takes care of those. But we still need to respect context cancellation and let UpdateUserTimestamp return early without waiting for the worker's response.

To achieve this we'll use Go's select statements for both sending the request to the worker and receiving the result. This way if the context is canceled, we return immediately with a "cancelled" error.

Important: From now on ReplyTo channels must be buffered. This prevents the worker from blocking if the caller's context is canceled and they don't read the result.

func UpdateUserTimestamp(ctx context.Context, userID int) error {
    // Prepare a request to the worker.
    // A ReplyTo channel is used by the worker to send us back a result.
    req := updateUserTimestampRequest{
        userID:  userID,
        ReplyTo: make(chan error, 1), // must be buffered
    }

    // Send request to the worker
    select {
    case <-ctx.Done():
        return ctx.Err()
    case updateUserTimestampQueue <- req:
    }

    // Block and wait for the result
    select {
    case <-ctx.Done():
        return ctx.Err()
    case err := <-req.ReplyTo:
        return err
    }
}

How to Batch Data Arriving in Real-Time?

Before we look at the worker implementation, let's think about how to batch data arriving in real time. Suppose we want to send queries to the database in batches of 10. We can just read from the queue until we have 10 items, then send them to the database. That will work fine, mostly...

What if it's a quiet day and UpdateUserTimestamp is called only once per minute? It would take 10 minutes to fill up the batch, and our database updates would be delayed by 10 minutes - making it the world's most patient real-time system.

To deal with this we need to introduce a timeout. If the batch isn't full after a certain time, we send it to the database anyway. The timeout can be low, like 10ms or even 1ms. This way during high traffic the batch fills up immediately, and during quiet periods we would introduce at most 1ms of latency — keeping the system as real-time as we want.

The Worker

To simplify the worker implementation, we'll use rill, a concurrency toolkit I’ve built. It provides utilities like rill.Batch, which groups incoming items into batches with a timeout (exactly as described earlier), and rill.ForEach, which helps control the level of concurrency when processing those batches. Thanks to these tools, the worker implementation becomes remarkably compact.

Here's the complete worker:

func updateUserTimestampWorker(batchSize int, batchTimeout time.Duration, concurrency int, dbTimeout time.Duration) {
    // Start with a stream of update requests
    requests := rill.FromChan(updateUserTimestampQueue, nil)

    // Group requests into batches with timeout
    requestBatches := rill.Batch(requests, batchSize, batchTimeout)

    // Process batches with controlled database concurrency
    _ = rill.ForEach(requestBatches, concurrency, func(batch []updateUserTimestampRequest) error {
        // Create a slice of user IDs
        ids := make([]int, len(batch))
        for i, req := range batch {
            ids[i] = req.userID
        }

        // Execute batched update
        dbCtx, cancel := context.WithTimeout(context.Background(), dbTimeout)
        defer cancel()
        err := sendQueryToDB(dbCtx,
            "UPDATE users SET last_active_at = NOW() WHERE id IN (?)",
            ids,
        )

        // Send result back to all callers in this batch
        for _, req := range batch {
            req.ReplyTo <- err
            close(req.ReplyTo)
        }
        return nil
    })
}

Putting it All Together

That's pretty much it. We've built a system that transparently batches real-time updates to the database. The developer just calls UpdateUserTimestamp as usual, and the system takes care of the rest. The worker reads user IDs from the queue, groups them into batches, and sends them to the database. The system respects context cancellation and handles errors properly.

And as a bonus, we can explicitly control the level of concurrency.

To see the solution in action, let's write a main function that spawns 100 goroutines each calling our UpdateUserTimestamp function with a different userID.

func main() {
    ctx := context.Background()

    // Start the worker
    go updateUserTimestampWorker(
        7,                   // Batch size
        10*time.Millisecond, // Batch timeout
        2,                   // Concurrency
        10*time.Second,      // Query timeout
    )

    // Simulate many concurrent goroutines calling UpdateUserTimestamp
    var wg sync.WaitGroup

    for i := 1; i <= 100; i++ {
        wg.Add(1)
        go func(userID int) {
            defer wg.Done()

            err := UpdateUserTimestamp(ctx, userID)
            if err != nil {
                fmt.Println("Error updating user timestamp:", err)
            }
        }(i)
    }

    wg.Wait()
}

// Simulate a database query
func sendQueryToDB(ctx context.Context, query string, args ...any) error {
    for _, arg := range args {
        query = strings.Replace(query, "?", fmt.Sprint(arg), 1)
    }
    fmt.Println("Executed:", query)
    return nil
}

The full code is available in the Go playground: https://goplay.tools/snippet/wTmEwFiYBvu

Conclusion

This pattern demonstrates how to combine batching for performance with real-time responsiveness using Go's concurrency primitives. By using channels, goroutines, and rill, we built a system that efficiently batches database updates while staying reactive.

The same approach works well for other scenarios where batching makes sense. Even SELECT queries can benefit from it, though the worker would need additional logic to route results back to callers. Not to mention a variety of non-real-time use cases where batching can be beneficial.

Of course, there are other ways to handle frequent updates, for example:

Store activity timestamps in Redis or another fast storage
Queue updates in an external queue (e.g RabbitMQ) and batch them in the consumer, making the pattern distributed
Update each user at most once per time window (e.g., 4 hours) and skip the rest of the updates.

This channel-based solution hits a sweet spot for systems that need performance of batching, while keeping the code simple and maintainable.

Originally published at destel.dev