2015-06-16 04:27:40 +08:00
package gracehttp_test
2012-06-05 06:21:10 +08:00
import (
2013-08-21 02:29:40 +08:00
"crypto/tls"
2012-06-05 06:21:10 +08:00
"encoding/json"
"flag"
2013-08-20 14:12:26 +08:00
"fmt"
2012-06-05 06:21:10 +08:00
"log"
"net/http"
"os"
2013-08-21 04:15:00 +08:00
"strings"
2013-08-20 14:12:26 +08:00
"sync"
2015-06-16 04:27:40 +08:00
"testing"
2012-06-05 06:21:10 +08:00
"time"
2013-08-20 14:12:26 +08:00
2024-04-18 14:09:13 +08:00
"grace/gracehttp"
2012-06-05 06:21:10 +08:00
)
2017-02-19 06:52:39 +08:00
const preStartProcessEnv = "GRACEHTTP_PRE_START_PROCESS"
2015-06-16 04:27:40 +08:00
func TestMain ( m * testing . M ) {
const (
testbinKey = "GRACEHTTP_TEST_BIN"
testbinValue = "1"
)
if os . Getenv ( testbinKey ) == testbinValue {
testbinMain ( )
return
}
if err := os . Setenv ( testbinKey , testbinValue ) ; err != nil {
panic ( err )
}
os . Exit ( m . Run ( ) )
}
2012-06-05 06:21:10 +08:00
type response struct {
Sleep time . Duration
Pid int
2013-08-21 04:15:00 +08:00
Error string ` json:",omitempty" `
2012-06-05 06:21:10 +08:00
}
2014-04-03 05:59:35 +08:00
// Wait for 10 consecutive responses from our own pid.
//
// This prevents flaky tests that arise from the fact that we have the
// perfectly acceptable (read: not a bug) condition where both the new and the
// old servers are accepting requests. In fact the amount of time both are
// accepting at the same time and the number of requests that flip flop between
// them is unbounded and in the hands of the various kernels our code tends to
// run on.
//
// In order to combat this, we wait for 10 successful responses from our own
// pid. This is a somewhat reliable way to ensure the old server isn't
// serving anymore.
2013-08-21 02:29:40 +08:00
func wait ( wg * sync . WaitGroup , url string ) {
2014-04-03 05:59:35 +08:00
var success int
2013-08-20 14:12:26 +08:00
defer wg . Done ( )
for {
2013-08-21 04:28:03 +08:00
res , err := http . Get ( url )
2013-08-21 02:29:40 +08:00
if err == nil {
2013-08-21 04:28:03 +08:00
// ensure it isn't a response from a previous instance
defer res . Body . Close ( )
var r response
if err := json . NewDecoder ( res . Body ) . Decode ( & r ) ; err != nil {
2013-09-19 08:58:04 +08:00
log . Fatalf ( "Error decoding json: %s" , err )
2013-08-21 04:28:03 +08:00
}
if r . Pid == os . Getpid ( ) {
2014-04-03 05:59:35 +08:00
success ++
if success == 10 {
return
}
continue
2013-08-21 04:28:03 +08:00
}
2013-08-21 02:29:40 +08:00
} else {
2014-04-03 05:59:35 +08:00
success = 0
2013-08-21 04:15:00 +08:00
// we expect connection refused
if ! strings . HasSuffix ( err . Error ( ) , "connection refused" ) {
e2 := json . NewEncoder ( os . Stderr ) . Encode ( & response {
Error : err . Error ( ) ,
Pid : os . Getpid ( ) ,
} )
if e2 != nil {
log . Fatalf ( "Error writing error json: %s" , e2 )
}
2013-08-21 02:29:40 +08:00
}
2013-08-20 14:12:26 +08:00
}
2013-08-20 13:51:32 +08:00
}
}
2013-08-21 02:29:40 +08:00
func httpsServer ( addr string ) * http . Server {
cert , err := tls . X509KeyPair ( localhostCert , localhostKey )
if err != nil {
2013-09-19 08:58:04 +08:00
log . Fatalf ( "error loading cert: %v" , err )
2013-08-21 02:29:40 +08:00
}
return & http . Server {
Addr : addr ,
Handler : newHandler ( ) ,
TLSConfig : & tls . Config {
NextProtos : [ ] string { "http/1.1" } ,
Certificates : [ ] tls . Certificate { cert } ,
} ,
}
}
2015-06-16 04:27:40 +08:00
func testbinMain ( ) {
2013-08-21 02:29:40 +08:00
var httpAddr , httpsAddr string
2017-02-19 06:52:39 +08:00
var testOption int
2013-08-21 02:29:40 +08:00
flag . StringVar ( & httpAddr , "http" , ":48560" , "http address to bind to" )
flag . StringVar ( & httpsAddr , "https" , ":48561" , "https address to bind to" )
2017-02-19 06:52:39 +08:00
flag . IntVar ( & testOption , "testOption" , - 1 , "which option to test on ServeWithOptions" )
2012-06-05 06:21:10 +08:00
flag . Parse ( )
2013-08-20 14:12:26 +08:00
2013-08-21 02:29:40 +08:00
// we have self signed certs
http . DefaultTransport = & http . Transport {
DisableKeepAlives : true ,
TLSClientConfig : & tls . Config {
InsecureSkipVerify : true ,
} ,
}
2013-08-20 14:12:26 +08:00
// print json to stderr once we can successfully connect to all three
// addresses. the ensures we only print the line once we're ready to serve.
go func ( ) {
var wg sync . WaitGroup
2013-08-21 02:29:40 +08:00
wg . Add ( 2 )
2014-04-03 05:59:35 +08:00
go wait ( & wg , fmt . Sprintf ( "http://%s/sleep/?duration=1ms" , httpAddr ) )
go wait ( & wg , fmt . Sprintf ( "https://%s/sleep/?duration=1ms" , httpsAddr ) )
2013-08-20 14:12:26 +08:00
wg . Wait ( )
2015-06-16 04:27:40 +08:00
err := json . NewEncoder ( os . Stderr ) . Encode ( & response { Pid : os . Getpid ( ) } )
2013-08-20 14:12:26 +08:00
if err != nil {
log . Fatalf ( "Error writing startup json: %s" , err )
}
} ( )
2013-08-20 13:51:32 +08:00
2017-02-19 06:52:39 +08:00
servers := [ ] * http . Server {
2013-08-21 02:29:40 +08:00
& http . Server { Addr : httpAddr , Handler : newHandler ( ) } ,
httpsServer ( httpsAddr ) ,
2017-02-19 06:52:39 +08:00
}
if testOption == - 1 {
err := gracehttp . Serve ( servers ... )
if err != nil {
log . Fatalf ( "Error in gracehttp.Serve: %s" , err )
}
} else {
if testOption == testPreStartProcess {
switch os . Getenv ( preStartProcessEnv ) {
case "" :
err := os . Setenv ( preStartProcessEnv , "READY" )
if err != nil {
log . Fatalf ( "testbin (first incarnation) could not set %v to 'ready': %v" , preStartProcessEnv , err )
}
case "FIRED" :
// all good, reset for next round
err := os . Setenv ( preStartProcessEnv , "READY" )
if err != nil {
log . Fatalf ( "testbin (second incarnation) could not reset %v to 'ready': %v" , preStartProcessEnv , err )
}
case "READY" :
log . Fatalf ( "failure to update startup hook before new process started" )
default :
log . Fatalf ( "something strange happened with %v: it ended up as %v, which is not '', 'FIRED', or 'READY'" , preStartProcessEnv , os . Getenv ( preStartProcessEnv ) )
}
err := gracehttp . ServeWithOptions (
servers ,
gracehttp . PreStartProcess ( func ( ) error {
err := os . Setenv ( preStartProcessEnv , "FIRED" )
if err != nil {
log . Fatalf ( "startup hook could not set %v to 'fired': %v" , preStartProcessEnv , err )
}
return nil
} ) ,
)
if err != nil {
log . Fatalf ( "Error in gracehttp.Serve: %s" , err )
}
}
2012-06-05 06:21:10 +08:00
}
}
func newHandler ( ) http . Handler {
mux := http . NewServeMux ( )
mux . HandleFunc ( "/sleep/" , func ( w http . ResponseWriter , r * http . Request ) {
duration , err := time . ParseDuration ( r . FormValue ( "duration" ) )
if err != nil {
http . Error ( w , err . Error ( ) , 400 )
}
time . Sleep ( duration )
err = json . NewEncoder ( w ) . Encode ( & response {
Sleep : duration ,
Pid : os . Getpid ( ) ,
} )
if err != nil {
log . Fatalf ( "Error encoding json: %s" , err )
}
} )
return mux
}
2013-08-21 02:29:40 +08:00
// localhostCert is a PEM-encoded TLS cert with SAN IPs
// "127.0.0.1" and "[::1]", expiring at the last second of 2049 (the end
// of ASN.1 time).
// generated from src/pkg/crypto/tls:
// go run generate_cert.go --rsa-bits 512 --host 127.0.0.1,::1,example.com --ca --start-date "Jan 1 00:00:00 1970" --duration=1000000h
var localhostCert = [ ] byte ( ` -- -- - BEGIN CERTIFICATE -- -- -
MIIBdzCCASOgAwIBAgIBADALBgkqhkiG9w0BAQUwEjEQMA4GA1UEChMHQWNtZSBD
bzAeFw03MDAxMDEwMDAwMDBaFw00OTEyMzEyMzU5NTlaMBIxEDAOBgNVBAoTB0Fj
bWUgQ28wWjALBgkqhkiG9w0BAQEDSwAwSAJBALyCfqwwip8BvTKgVKGdmjZTU8DD
ndR + WALmFPIRqn89bOU3s30olKiqYEju / SFoEvMyFRT / TWEhXHDaufThqaMCAwEA
AaNoMGYwDgYDVR0PAQH / BAQDAgCkMBMGA1UdJQQMMAoGCCsGAQUFBwMBMA8GA1Ud
EwEB / wQFMAMBAf8wLgYDVR0RBCcwJYILZXhhbXBsZS5jb22HBH8AAAGHEAAAAAAA
AAAAAAAAAAAAAAEwCwYJKoZIhvcNAQEFA0EAr / 0 9 uy108p51rheIOSnz4zgduyTl
M + 4 AmRo8 / U1twEZLgfAGG / GZjREv2y4mCEUIM3HebCAqlA5jpRg76Rf8jw ==
-- -- - END CERTIFICATE -- -- - ` )
// localhostKey is the private key for localhostCert.
var localhostKey = [ ] byte ( ` -- -- - BEGIN RSA PRIVATE KEY -- -- -
MIIBOQIBAAJBALyCfqwwip8BvTKgVKGdmjZTU8DDndR + WALmFPIRqn89bOU3s30o
lKiqYEju / SFoEvMyFRT / TWEhXHDaufThqaMCAwEAAQJAPXuWUxTV8XyAt8VhNQER
LgzJcUKb9JVsoS1nwXgPksXnPDKnL9ax8VERrdNr + nZbj2Q9cDSXBUovfdtehcdP
qQIhAO48ZsPylbTrmtjDEKiHT2Ik04rLotZYS2U873J6I7WlAiEAypDjYxXyafv /
Yo1pm9onwcetQKMW8CS3AjuV9Axzj6cCIEx2Il19fEMG4zny0WPlmbrcKvD / DpJQ
4 FHrzsYlIVTpAiAas7S1uAvneqd0l02HlN9OxQKKlbUNXNme + rnOnOGS2wIgS0jW
zl1jvrOSJeP1PpAHohWz6LOhEr8uvltWkN6x3vE =
-- -- - END RSA PRIVATE KEY -- -- - ` )