Simulate a tcp connection in Go - networking

In Go, a TCP connection (net.Conn) is a io.ReadWriteCloser. I'd like to test my network code by simulating a TCP connection. There are two requirements that I have:
the data to be read is stored in a string
whenever data is written, I'd like it to be stored in some kind of buffer which I can access later
Is there a data structure for this, or an easy way to make one?

No idea if this existed when the question was asked, but you probably want net.Pipe() which provides you with two full duplex net.Conn instances which are linked to each other

EDIT: I've rolled this answer into a package which makes things a bit simpler - see here: https://github.com/jordwest/mock-conn
While Ivan's solution will work for simple cases, keep in mind that a real TCP connection is actually two buffers, or rather pipes. For example:
Server | Client
---------+---------
reads <=== writes
writes ===> reads
If you use a single buffer that the server both reads from and writes to, you could end up with the server talking to itself.
Here's a solution that allows you to pass a MockConn type as a ReadWriteCloser to the server. The Read, Write and Close functions simply proxy through to the functions on the server's end of the pipes.
type MockConn struct {
ServerReader *io.PipeReader
ServerWriter *io.PipeWriter
ClientReader *io.PipeReader
ClientWriter *io.PipeWriter
}
func (c MockConn) Close() error {
if err := c.ServerWriter.Close(); err != nil {
return err
}
if err := c.ServerReader.Close(); err != nil {
return err
}
return nil
}
func (c MockConn) Read(data []byte) (n int, err error) { return c.ServerReader.Read(data) }
func (c MockConn) Write(data []byte) (n int, err error) { return c.ServerWriter.Write(data) }
func NewMockConn() MockConn {
serverRead, clientWrite := io.Pipe()
clientRead, serverWrite := io.Pipe()
return MockConn{
ServerReader: serverRead,
ServerWriter: serverWrite,
ClientReader: clientRead,
ClientWriter: clientWrite,
}
}
When mocking a 'server' connection, simply pass the MockConn in place of where you would use the net.Conn (this obviously implements the ReadWriteCloser interface only, you could easily add dummy methods for LocalAddr() etc if you need to support the full net.Conn interface)
In your tests you can act as the client by reading and writing to the ClientReader and ClientWriter fields as needed:
func TestTalkToServer(t *testing.T) {
/*
* Assumes that NewMockConn has already been called and
* the server is waiting for incoming data
*/
// Send a message to the server
fmt.Fprintf(mockConn.ClientWriter, "Hello from client!\n")
// Wait for the response from the server
rd := bufio.NewReader(mockConn.ClientReader)
line, err := rd.ReadString('\n')
if line != "Hello from server!" {
t.Errorf("Server response not as expected: %s\n", line)
}
}

Why not using bytes.Buffer? It's an io.ReadWriter and has a String method to get the stored data. If you need to make it an io.ReadWriteCloser, you could define you own type:
type CloseableBuffer struct {
bytes.Buffer
}
and define a Close method:
func (b *CloseableBuffer) Close() error {
return nil
}

In majority of the cases you do not need to mock net.Conn.
You only have to mock stuff that will add time to your tests, prevent tests from running in parallel (using shared resources like the hardcoded file name) or can lead to outages (you can potentially exhaust the connection limit or ports but in most of the cases it is not a concern, when you run your tests in isolation).
Not mocking has an advantage of more precise testing of what you want to test with a real thing.
https://www.accenture.com/us-en/blogs/software-engineering-blog/to-mock-or-not-to-mock-is-that-even-a-question
Instead of mocking net.Conn, you can write a mock server, run it in a goroutine in your test and connect to it using real net.Conn
A quick and dirty example:
port := someRandomPort()
srv := &http.Server{Addr: port}
go func(msg string) {
http.HandleFunc("/hello", myHandleFUnc)
srv.ListenAndServe()
}
myTestCodeUsingConn(port)
srv.Shutdown(context.TODO())

Related

Go HTTP RoundTripper: Preventing Connection Reuse Based on Response

I have a use case where I want to use an HTTP client in Go with pooled connections (connection re-use), but with the special case where a connection is intentionally closed (not allowed for re-use) if a request on that connection returns a specific HTTP status code.
I've implemented a custom http.RoundTripper, which wraps an http.Transport, and can inspect the response status code. However, I can't seem to find a way to prevent the http.Transport from re-using that connection, without also preventing it from re-using any other connection.
Is this possible using the net/http package? If not, any suggested workaround for accomplishing this?
My current code looks something like this:
type MyTransport struct {
transport *http.Transport
}
func (mt *MyTransport) RoundTrip(req *http.Request) (*http.Response, error) {
resp, err := tt.transport.RoundTrip(req)
if err != nil {
return resp, err
}
if resp.StatusCode == 567 {
// HERE:
// Do something to prevent re-use of this connection
}
return resp, err
}

How to interrupt an HTTP handler?

Say I have a http handler like this:
func ReallyLongFunction(w http.ResponseWriter, r *http.Request) {
fmt.Fprintf(w, "Hello World!")
// run code that takes a long time here
// Executing dd command with cmd.Exec..., etc.
})
Is there a way I can interrupt this function if the user refreshes the page or kills the request some other way without running the subsequent code and how would I do it?
I tried doing this:
notify := r.Context().Done()
go func() {
<-notify
println("Client closed the connection")
s.downloadCleanup()
return
}()
but the code after whenever I interrupt it still runs anyway.
There's no way to forcibly tear a goroutine down from any code external to that goroutine.
Hence the only way to actually interrupt processing is to periodically check whether the client is gone (or whether there's another signal to stop processing).
Basically that would amount to structuring your handler something like this
func ReallyLongFunction(w http.ResponseWriter, r *http.Request) {
fmt.Fprintf(w, "Hello World!")
done := r.Context().Done()
// Check wheteher we're done
// do some small piece of stuff
// check whether we're done
// do another small piece of stuff
// …rinse, repeat
})
Now a way to check whether there was something written to a channel, but without blocking the operation is to use the "select with default" idiom:
select {
case <- done:
// We're done
default:
}
This statemept executes the code in the "// We're done" block if and only if done was written to or was closed (which is the case with contexts), and otherwis the empty block in the default branch is executed.
So we can refactor that to something like
func ReallyLongFunction(w http.ResponseWriter, r *http.Request) {
fmt.Fprintf(w, "Hello World!")
done := r.Context().Done()
closed := func () bool {
select {
case <- done:
return true
default:
return false
}
}
if closed() {
return
}
// do some small piece of stuff
if closed() {
return
}
// do another small piece of stuff
// …rinse, repeat
})
Stopping an external process started in an HTTP handler
To address the OP's comment…
The os/exec.Cmd type has the Process field, which is of type os.Process and that type supports the Kill method which forcibly brings the running process down.
The only problem is that exec.Cmd.Run blocks until the process exits,
so the goroutine which is executing it cannot execute other code, and if exec.Cmd.Run is called in an HTTP handler, there's no way to cancel it.
How to best handle running a program in such an asynchronous manner heavily depends on how the process itself is organized but I'd roll like this:
In the handler, prepare the process and then start it using exec.Cmd.Start (as opposed to Run).
Check the error value Start have returned: if it's nil
the process has managed to start OK. Otherwise somehow communicate the failure to the client and quit the handler.
Once the process is known to had started, the exec.Cmd value
has some of its fields populated with process-related information;
of particular interest is the Process field which is of type
os.Process: that type has the Kill method which may be used to forcibly bring the process down.
Start a goroutine and pass it that exec.Cmd value and a channel of some suitable type (see below).
That goroutine should call Wait on it and once it returns,
it should communicate that fact back to the originating goroutine over that channel.
Exactly what to communicate, is an open question as it depends
on whether you want to collect what the process wrote to its standard
output and error streams and/or may be some other data related to the process' activity.
After sending the data, that goroutine exits.
The main goroutine (executing the handler) should just call exec.Cmd.Process.Kill when it detect the handler should terminate.
Killing the process eventually unblocks the goroutine which is executing Wait on that same exec.Cmd value as the process exits.
After killing the process, the handler goroutine waits on the channel to hear back from the goroutine watching the process. The handler does something with that data (may be logs it or whatever) and exits.
You should cancel the goroutine from inside, so for a long calculation task, you may provide checkpoints, to stop and check for the cancelation:
Here is the tested code for the server which has e.g. long calculation task and checkpoints for the cancelation:
package main
import (
"fmt"
"io"
"log"
"net/http"
"time"
)
func main() {
http.HandleFunc(`/`, func(w http.ResponseWriter, r *http.Request) {
ctx := r.Context()
log.Println("wait a couple of seconds ...")
for i := 0; i < 10; i++ { // long calculation
select {
case <-ctx.Done():
log.Println("Client closed the connection:", ctx.Err())
return
default:
fmt.Print(".")
time.Sleep(200 * time.Millisecond) // long calculation
}
}
io.WriteString(w, `Hi`)
log.Println("Done.")
})
log.Println(http.ListenAndServe(":8081", nil))
}
Here is the client code, which times out:
package main
import (
"io/ioutil"
"log"
"net/http"
"time"
)
func main() {
log.Println("HTTP GET")
client := &http.Client{
Timeout: 1 * time.Second,
}
r, err := client.Get(`http://127.0.0.1:8081/`)
if err != nil {
log.Fatal(err)
}
defer r.Body.Close()
bs, err := ioutil.ReadAll(r.Body)
if err != nil {
log.Fatal(err)
}
log.Println("HTTP Done.")
log.Println(string(bs))
}
You may use normal browser to check for not canclation, or close it, refresh it , disconect it, or ..., for the cancelation.

Async work after response

I am trying to implement http server that:
Calculate farther redirect using some logic
Redirect user
Log user data
The goal is to achieve maximum throughput (at least 15k rps). In order to do this, I want to save log asynchronously. I'm using kafka as logging system and separate logging block of code into separate goroutine. Overall example of current implementation:
package main
import (
"github.com/confluentinc/confluent-kafka-go/kafka"
"net/http"
"time"
"encoding/json"
)
type log struct {
RuntimeParam string `json:"runtime_param"`
AsyncParam string `json:"async_param"`
RemoteAddress string `json:"remote_address"`
}
var (
producer, _ = kafka.NewProducer(&kafka.ConfigMap{
"bootstrap.servers": "localhost:9092,localhost:9093",
"queue.buffering.max.ms": 1 * 1000,
"go.delivery.reports": false,
"client.id": 1,
})
topicName = "log"
)
func main() {
siteMux := http.NewServeMux()
siteMux.HandleFunc("/", httpHandler)
srv := &http.Server{
Addr: ":8080",
Handler: siteMux,
ReadTimeout: 2 * time.Second,
WriteTimeout: 5 * time.Second,
IdleTimeout: 10 * time.Second,
}
if err := srv.ListenAndServe(); err != nil {
panic(err)
}
}
func httpHandler(w http.ResponseWriter, r *http.Request) {
handlerLog := new(log)
handlerLog.RuntimeParam = "runtimeDataString"
http.Redirect(w, r, "http://google.com", 301)
go func(goroutineLog *log, request *http.Request) {
goroutineLog.AsyncParam = "asyncDataString"
goroutineLog.RemoteAddress = r.RemoteAddr
jsonLog, err := json.Marshal(goroutineLog)
if err == nil {
producer.ProduceChannel() <- &kafka.Message{
TopicPartition: kafka.TopicPartition{Topic: &topicName, Partition: kafka.PartitionAny},
Value: jsonLog,
}
}
}(handlerLog, r)
}
The questions are:
Is it correct/efficient to use separate goroutine to implement async logging or should I use a different approach? (workers and channels for example)
Maybe there is a way to further improve performance of server, that I'm missing?
Yes, this is correct and efficient use of a goroutine (as Flimzy pointed in the comments). I couldn't agree more, this is a good approach.
The problem is that the handler may finish executing before the goroutine started processing everything and the request (which is a pointer) may be gone or you may have some races down the middleware stack. I read your comments, that it isn't your case, but in general, you shouldn't pass a request to a goroutine. As I can see from your code, you're really using only RemoteAddr from the request and why not to redirect straight away and put logging in the defer statement? So, I'd rewrite your handler a bit:
func httpHandler(w http.ResponseWriter, r *http.Request) {
http.Redirect(w, r, "http://google.com", 301)
defer func(runtimeDataString, RemoteAddr string) {
handlerLog := new(log)
handlerLog.RuntimeParam = runtimeDataString
handlerLog.AsyncParam = "asyncDataString"
handlerLog.RemoteAddress = RemoteAddr
jsonLog, err := json.Marshal(handlerLog)
if err == nil {
producer.ProduceChannel() <- &kafka.Message{
TopicPartition: kafka.TopicPartition{Topic: &topicName, Partition: kafka.PartitionAny},
Value: jsonLog,
}
}
}("runtimeDataString", r.RemoteAddr)
}
The goroutines unlikely improve performance of your server as you just send the response earlier and those kafka connections could pile up in the background and slow down the whole server. If you find this as the bottleneck, you may consider saving logs locally and sending them to kafka in another process (or pool of workers) outside of your server. This may spread the workload over time (like sending fewer logs when you have more requests and vice versa).

How come the redis-benchmark command is not following the redis protocol?

I was reading in directly from a tcp connection after running the redis-benchmark command and as far as I can tell, redis-benchmark is NOT following the redis protocol.
The redis protocol is as stated in its website:
The way RESP is used in Redis as a request-response protocol is the
following:
Clients send commands to a Redis server as a RESP Array of Bulk Strings.
The server replies with one of the RESP types according to the command implementation.
Meaning that a correct client implementation must always send RESP arrays of bulk strings.
If that is true, then, anything that does not start with a * is considered a syntax error (since its not an RESP array).
Thus, if one were to send a ping command to a redis-server, then it must be sent as a resp array of length 1 with 1 bulk string containing the word ping. For example:
"*1\r\n$4\r\nPING\r\n"
However, whenever I listen directly to the redis-benchmark command and read its tcp connection I get instead:
"PING\r\n"
which does not follow the redis protocol. Is that a bug or is there something implied in the redis protocol that makes pings special? As far as I could tell I couldn't find anything that said that pings were special, nor that length 1 commands were special. Does someone know whats going on?
To see reproduce these results yourself you can copy my code to inspect it directly:
package main
import (
"fmt"
"log"
"net"
)
func main() {
RedisBenchmark()
}
func RedisBenchmark() {
url := "127.0.0.1:6379"
fmt.Println("listen: ", url)
ln, err := net.Listen("tcp", url) //announces on local network
if err != nil {
log.Fatal(err)
}
for {
conn, err := ln.Accept() //waits and returns the next connection to the listener
if err != nil {
log.Fatal(err)
}
tcpConn := conn.(*net.TCPConn)
go HandleConnection(tcpConn)
}
}
func HandleConnection(tcpConn *net.TCPConn) {
b := make([]byte, 256) //TODO how much should I read at a time?
n, err := tcpConn.Read(b)
if err != nil {
fmt.Println("n: ", n)
log.Fatal(err)
}
fmt.Printf("+++++> raw input string(b): %q\n", string(b))
msg := string(b[:n])
fmt.Printf("+++++> raw input msg: %q\n", msg)
}
and run it using go with:
go run main.go
followed on a different terminal (or tmux pane):
redis-benchmark
for all the test or if you only want to run ping with 1 client:
redis-benchmark -c 1 -t ping -n 1
you can see the details of how I am running it with the flags at: http://redis.io/topics/benchmarks
That is called an inline command. Check the Inline Commands section of the Redis Protocol article.
You can refer to the source code to find out the differences between inline command and RESP.
readQueryFromClient
|--> if command begins with * --> processInlineBuffer()process it as RESP
|
|--> if command not begins with * --> processMultibulkBuffer():process it as inline command
RESP is a more efficent way to parse the command for the Redis Server

Inspecting the body of an HTTP request with gocraft Middleware

I've been using the gocraft-web package so far to do some development on an HTTP service. It's really great because you can stick middleware in it to check for stuff like the presence of a Cookie in the header.
At the moment I am wanting to implement request signing. Getting the client to sign the request is easy enough, but I am wanting to check it for all endpoints with a common piece of middleware. Basically the middleware needs to find the key to check against, compute the request HMAC, and check it against the supplied HMAC (presumably in the Authorization Header).
Computing the actual HMAC is really easy in go.
The problem is: reading the message in middleware makes it unavailable to the final endpoint.
The best solution I have come up with (example shown below) is to read everything from the Request in the middleware and stuffing it back into a bytes.Buffer for later reading. Is there a better way to do this? The current implementation seems a bit hackish.
Reading everything into memory sucks, but I can probably just put my service behind a proxy and limit the size of requests anyways. The actual content will always be pretty small(under 5 kilobytes). The extra copy introduced by this approach is likely to be quite slow, but computing the HMAC of a message is not exactly cheap to begin with.
The advantage to this is that it is transparent: it will work with any other go http code that just expects to read from Request.Body without any magic.
I suppose I could be a bit slicker and use a io.TeeReader.
This is my solution so far. If you post to localhost:3300 some JSON it prints the sha512 to the terminal in the server process, but also the response is able to contain a listing of the keys & values in it.
package main
import "fmt"
import "github.com/gocraft/web"
import "net/http"
import "bytes"
import "crypto/sha512"
import "io"
import "encoding/hex"
import "encoding/json"
type Context struct{}
type echoer struct {
*bytes.Buffer
}
func (e echoer) Close() error {
//Just do nothing to make the interface happy
return nil
}
func middlewareThatLooksAtBody(rw web.ResponseWriter, req *web.Request, next web.NextMiddlewareFunc) {
var replacement echoer
replacement.Buffer = &bytes.Buffer{}
hash := sha512.New()
hash.Write([]byte(req.Method))
reader := req.Body
var bytes []byte = make([]byte, 64)
for {
amount, err := reader.Read(bytes)
fmt.Printf("Read %d bytes\n", amount)
if err != nil {
if err == io.EOF {
break
}
panic(err)
}
if amount == 0 {
break
}
hash.Write(bytes)
replacement.Write(bytes)
}
//Is this needed?
reader.Close()
//replacement.Seek(0, 0)
req.Body = replacement
fmt.Printf("%v\n", hex.EncodeToString(hash.Sum(nil)))
next(rw, req)
}
func echoJson(rw web.ResponseWriter, req *web.Request) {
dec := json.NewDecoder(req.Body)
var obj map[string]interface{}
err := dec.Decode(&obj)
if err != nil {
rw.WriteHeader(http.StatusBadRequest)
fmt.Fprintf(rw, "%v\n", err)
return
}
for k, v := range obj {
fmt.Fprintf(rw, "%v = %v\n", k, v)
}
}
func main() {
router := web.New(Context{})
router.Middleware(middlewareThatLooksAtBody)
router.Post("/", echoJson)
http.ListenAndServe("localhost:3300", router)
}
From your description, it looks like you need to read all the bytes from the request body, regardless of what your handlers will do.
If so, then you have at least a couple of options that would avoid the extra copy:
1) Store the read contents inside your gocraft context.
2) Do all body data processing and validation in the middleware and store the results of the processing in the context.
Granted, this means that your handlers now must know that they should look for the contents in the context instead of the req.Body.
I think it's a decent trade-off though, given your requirements.

Resources