Start Go HTTP server, do something, then shut it down once task is done

Question

I'm setting up an OAuth2 flow from a CLI application, I'm working on. I need to create a temporary HTTP server for the provider to send the callback to, e.g. localhost:8080/callback

Once the provider has sent the details I need, I want to be able to shut the HTTP server down, just to keep everything clean. I think what I'm looking for is Routines and Wait Groups, but I'm still quite new to this area.

This is what I have so far. I have redacted the part that sends the user to the provider, as my main issue is simply how to shut down the HTTP server once the token variable has been captured.

1. Server starts
2. User is directed to authorization URL at the provider site
3. User approves the request
4. Provider directs user back to localhost:8080/callback
5. URL includes client-side only params so I have to server HTML to use JS to capture the values and send it back to the server
6. Server receives token and can then shutdown

package main

import (
    "fmt"
    "log"
    "net/http"
    "sync"
)

func main() {
    // Start local HTTP serevr to listen for response
    serverDone := &sync.WaitGroup{}
    serverDone.Add(1)
    Start(serverDone)

    // ... Process to start OAuth2 flow
    // User is directed to provider website
    // User approves
    // Provider direct user back to localhost/callback
    
    serverDone.Wait()
}

func Start(wg *sync.WaitGroup) {

    srv := &http.Server{Addr: ":8080"}
    http.HandleFunc("/callback", func(w http.ResponseWriter, r *http.Request) {
        token := r.URL.Query().Get("token")
        if token != "" {
            fmt.Println("Found Token:", token)
            // Shut down server here
        } else {
            // Server HTML page to fetch token and return to server at /callback
        }
    })

    go func() {
        // let main know we are done cleaning up
        defer wg.Done()

        // ErrServerClosed on graceful close
        if err := srv.ListenAndServe(); err != http.ErrServerClosed {
            log.Fatalf("ListenAndServe(): %v", err)
        }
    }()
}

Answer 1

Use:

var ctxShutdown, cancel = context.WithCancel(context.Background())

Then:

cancel() // to say sorry, above.
// graceful-shutdown
err := srv.Shutdown(context.Background())

Try this:

package main

import (
    "context"
    "fmt"
    "log"
    "net/http"
    "sync"
)

func main() {
    serverDone := &sync.WaitGroup{}
    serverDone.Add(1)
    Start(serverDone)
    serverDone.Wait()
    fmt.Println("Done that.")
}

var ctxShutdown, cancel = context.WithCancel(context.Background())

func Start(wg *sync.WaitGroup) {
    srv := &http.Server{Addr: ":8080"}
    http.HandleFunc("/callback", func(w http.ResponseWriter, r *http.Request) {
        select {
        case <-ctxShutdown.Done():
            fmt.Println("Sorry: Shuting down ...")
            return
        default:
        }
        token := r.URL.Query().Get("token")
        if token != "" {
            fmt.Println("Found Token:", token)
            fmt.Println("Shuting down ...")
            // Shut down server here
            cancel() // to say sorry, above.
            // graceful-shutdown
            err := srv.Shutdown(context.Background())
            if err != nil {
                log.Println("server.Shutdown:", err)
            }

        } else {
            fmt.Fprintln(w, "Hi") // Server HTML page to fetch token and return to server at /callback
        }
    })

    go func() {
        defer wg.Done()
        if err := srv.ListenAndServe(); err != http.ErrServerClosed {
            log.Fatalf("ListenAndServe(): %v", err)
        }
        fmt.Println("Bye.")
    }()
}

Run and open http://127.0.0.1:8080/callback?token=2

Output:

Found Token: 2
Shuting down ...
Bye.
Done that.