How to write data to client and then exit process in golang http server?

Question

I'm writing an http server with self-update feature, the steps are:

1. client post to the server with the new executable, as the http body
2. server replace itself with the bytes
3. server echos a string "ok" to the client
4. server exit process
5. server will be restarted by systemd service

The problem is when the server calls os.Exit(), the client will receive EOF instead of "ok", my code looks like:

Server:

package main

import (
    "io"
    "net/http"
    "os"
)

func replace_self_executable(executable_bytes []byte) {
    // ...
}

func handle_self_update(w http.ResponseWriter, r *http.Request) {
    executable_bytes, _ := io.ReadAll(r.Body)

    replace_self_executable(executable_bytes)

    w.Write([]byte(`ok`))

    os.Exit(0)

    // other stuff
}

func main() {
    http.HandleFunc("/self_update", handle_self_update)

    http.ListenAndServe(":123", nil)
}

Client:

package main

import (
    "bytes"
    "fmt"
    "io"
    "net/http"
)

func main() {

    new_server_binary := []byte{ /*...*/ }

    resp, err := http.Post(`http://localhost:123/self_update`, ``, bytes.NewReader(new_server_binary))
    if err != nil {
        panic(err)
    }

    result, err := io.ReadAll(resp.Body)
    if err != nil {
        panic(err)
    }
    fmt.Println(`result`, string(result))
}

I also tried flushing and close client first, seems no difference:

    if f, ok := w.(http.Flusher); ok {
        f.Flush()
    }
    r.Body.Close()

How to send the "ok" to client?

Answer 1

Signal (e.g. via channel) to a routine outside the handler so that you can return from the request handler and gracefully shutdown the server:

func handle_self_update(w http.ResponseWriter, r *http.Request) {
    executable_bytes, _ := io.ReadAll(r.Body)

    replace_self_executable(executable_bytes)

    quitChan <- true

    w.Write([]byte(`ok`))
}

func main() {
    go func() {
        _ := <- quitChan
        // Graceful shutdown
    }

    // Start server
}

On an unrelated note - be very very very very careful with your implementation. Allowing a client to post you a new binary which you'll run on the server is a gigantic security hole and should have layers of authentication, signature verification, checksum validation, etc. Or better yet, let the client signal that an update is available, and have the server reach out to a known-good source for the updated binary (and still confirm signature & checksum).