homework-jianmu/tools/keeper/api/https_test.go

128 lines
3.5 KiB
Go

package api
import (
"context"
"crypto/ecdsa"
"crypto/elliptic"
crand "crypto/rand"
"crypto/tls"
"crypto/x509"
"crypto/x509/pkix"
"encoding/pem"
"fmt"
"log"
"math/big"
"net/http"
"net/http/httputil"
"net/url"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/taosdata/taoskeeper/db"
"github.com/taosdata/taoskeeper/util"
)
func TestHttps(t *testing.T) {
server := startProxy()
defer server.Shutdown(context.Background())
cfg := util.GetCfg()
cfg.TDengine.Usessl = true
cfg.TDengine.Port = 34443
CreateDatabase(cfg.TDengine.Username, cfg.TDengine.Password, cfg.TDengine.Host, cfg.TDengine.Port, cfg.TDengine.Usessl, cfg.Metrics.Database.Name, cfg.Metrics.Database.Options)
conn, err := db.NewConnectorWithDb(cfg.TDengine.Username, cfg.TDengine.Password, cfg.TDengine.Host, cfg.TDengine.Port, cfg.Metrics.Database.Name, cfg.TDengine.Usessl)
assert.NoError(t, err)
defer func() {
_, _ = conn.Query(context.Background(), fmt.Sprintf("drop database if exists %s", cfg.Metrics.Database.Name), util.GetQidOwn())
}()
data, err := conn.Query(context.Background(), "select server_version()", util.GetQidOwn())
assert.NoError(t, err)
assert.Equal(t, 1, len(data.Data))
}
func generateSelfSignedCert() (tls.Certificate, error) {
priv, err := ecdsa.GenerateKey(elliptic.P384(), crand.Reader)
if err != nil {
return tls.Certificate{}, err
}
notBefore := time.Now()
notAfter := notBefore.Add(365 * 24 * time.Hour)
serialNumber, err := crand.Int(crand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
if err != nil {
return tls.Certificate{}, err
}
template := x509.Certificate{
SerialNumber: serialNumber,
Subject: pkix.Name{
Organization: []string{"Your Company"},
},
NotBefore: notBefore,
NotAfter: notAfter,
KeyUsage: x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
BasicConstraintsValid: true,
}
certDER, err := x509.CreateCertificate(crand.Reader, &template, &template, &priv.PublicKey, priv)
if err != nil {
return tls.Certificate{}, err
}
certPEM := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: certDER})
keyPEM, err := x509.MarshalECPrivateKey(priv)
if err != nil {
return tls.Certificate{}, err
}
keyPEMBlock := pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: keyPEM})
return tls.X509KeyPair(certPEM, keyPEMBlock)
}
func startProxy() *http.Server {
// Generate self-signed certificate
cert, err := generateSelfSignedCert()
if err != nil {
log.Fatalf("Failed to generate self-signed certificate: %v", err)
}
target := "http://127.0.0.1:6041"
proxyURL, err := url.Parse(target)
if err != nil {
log.Fatalf("Failed to parse target URL: %v", err)
}
proxy := httputil.NewSingleHostReverseProxy(proxyURL)
proxy.ErrorHandler = func(w http.ResponseWriter, r *http.Request, e error) {
http.Error(w, "Proxy error", http.StatusBadGateway)
}
mux := http.NewServeMux()
mux.Handle("/", proxy)
server := &http.Server{
Addr: ":34443",
Handler: mux,
TLSConfig: &tls.Config{Certificates: []tls.Certificate{cert}},
// Setup server timeouts for better handling of idle connections and slowloris attacks
WriteTimeout: 10 * time.Second,
ReadTimeout: 10 * time.Second,
IdleTimeout: 30 * time.Second,
}
log.Println("Starting server on :34443")
go func() {
err = server.ListenAndServeTLS("", "")
if err != nil && err != http.ErrServerClosed {
log.Fatalf("Failed to start HTTPS server: %v", err)
}
}()
return server
}