如何有效地实现json tcp服务器并防止套接字泛滥？

Question

I am searching for most efficient solution, there are a lot of ways to read data from socket and decode json. I obviously should use json.Encoder and json.Decoder, because they are suitable for streaming nature of socket, but I have specific rule to prevent socket flooding, I must close connection if there is a single message > than 5 Kb. My message structure is JSON RPC.

In the following example I can check length and apply policy:

connbuf := bufio.NewReader(conn)
msg, err := connbuf.ReadBytes('
')
if len(msg) > 5 * 1024 {
    conn.Close()
}
...
var req JSONRequest
err = json.Unmarshal(message, &req)
...

But if client pushes megabytes of data without delimiter, this data will be in application, in msg variable already before server will disconnect client. Pretty vulnerable.

Second example uses Decoder, there is no chance to check size at all.

dec = json.NewDecoder(conn)
for {
    var req JSONRequest
    if err := dec.Decode(&req); err == io.EOF {
        break
    } else if err != nil {
        log.Println(err.Error())
        return err
   }
   ...
}

What is the best approach you can suggest to me? Thanks.

Answer 1

For the first example you can use ReadLine:

connbuff := bufio.NewReaderSize(conn, 5*1024)
msg, isPrefix, err := connbuff.ReadLine()
if isPrefix {
  // too long
}
...

If isPrefix is true then the line was too long. If you used a bufio.Scanner it actually already has a max token size of 64kb.

As Tim Cooper & Dave C said you can use io.LimitedReader for the second case, but there's one gotcha with the json decoder. It uses buffered IO, so it will read past the first request.

To fix that use a combination of io.MultiReader and io.LimitReader:

// to start with we have nothing buffered (an empty byte slice)
var buffered io.Reader = bytes.NewReader([]byte{})
for {
    // combine whatever was in buffered with conn, but only up to 5kb
    dec := json.NewDecoder(io.LimitReader(io.MultiReader(buffered, conn), 5*1024))
    var req string
    err := dec.Decode(&req)
    if err == io.EOF {
        break
    } else if err != nil {
        log.Fatalln(err)
    }
    // we probably read past the message, so save that to buffered
    buffered = dec.Buffered()
}