Advertisement
cooperlees

Untitled

Feb 14th, 2017
565
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 2.80 KB | None | 0 0
  1. diff --git a/cmd/chihaya/main.go b/cmd/chihaya/main.go
  2. index 9e48dad..7da79e1 100644
  3. --- a/cmd/chihaya/main.go
  4. +++ b/cmd/chihaya/main.go
  5. @@ -4,8 +4,14 @@
  6.  
  7. package main
  8.  
  9. -import "github.com/chihaya/chihaya"
  10. +import (
  11. + "github.com/chihaya/chihaya"
  12. + "github.com/chihaya/chihaya/config"
  13. + "os"
  14. +)
  15.  
  16. func main() {
  17. + config.DefaultConfig.TLSKeyPath = os.Getenv("THRIFT_TLS_CL_KEY_PATH")
  18. + config.DefaultConfig.TLSCertPath = os.Getenv("THRIFT_TLS_CL_CERT_PATH")
  19. chihaya.Boot()
  20. }
  21. diff --git a/config/config.go b/config/config.go
  22. index c46cb04..1297bf9 100644
  23. --- a/config/config.go
  24. +++ b/config/config.go
  25. @@ -87,6 +87,8 @@ type Config struct {
  26.  
  27. ClientWhitelistEnabled bool `json:"client_whitelist_enabled"`
  28. ClientWhitelist []string `json:"client_whitelist,omitempty"`
  29. + TLSCertPath string `json:"tls_cert_path"`
  30. + TLSKeyPath string `json:"tls_key_path"`
  31.  
  32. StatsConfig
  33. NetConfig
  34. @@ -135,6 +137,8 @@ var DefaultConfig = Config{
  35. },
  36.  
  37. ClientWhitelistEnabled: false,
  38. + TLSCertPath: "",
  39. + TLSKeyPath: "",
  40. }
  41.  
  42. // Open is a shortcut to open a file, read it, and generate a Config.
  43. diff --git a/http/http.go b/http/http.go
  44. index 31774b0..d907a17 100644
  45. --- a/http/http.go
  46. +++ b/http/http.go
  47. @@ -6,13 +6,18 @@
  48. package http
  49.  
  50. import (
  51. + "crypto/rand"
  52. + "crypto/tls"
  53. +
  54. "net"
  55. "net/http"
  56. "time"
  57. + "strings"
  58.  
  59. "github.com/golang/glog"
  60. "github.com/julienschmidt/httprouter"
  61. - "github.com/stretchr/graceful"
  62. + "github.com/tylerb/graceful"
  63. + "github.com/soheilhy/cmux"
  64.  
  65. "github.com/chihaya/chihaya/config"
  66. "github.com/chihaya/chihaya/stats"
  67. @@ -133,13 +138,47 @@ func Serve(cfg *config.Config, tkr *tracker.Tracker) {
  68. },
  69. }
  70.  
  71. + l, err := net.Listen("tcp", cfg.Addr)
  72. + if err != nil {
  73. + panic(err)
  74. + }
  75. +
  76. + // Create a cmux.
  77. + m := cmux.New(l)
  78. + httpl := m.Match(cmux.HTTP1Fast())
  79. + go grace.Serve(httpl)
  80. +
  81. + if cfg.TLSCertPath != "" && cfg.TLSKeyPath != "" {
  82. + tlsl := m.Match(cmux.Any())
  83. +
  84. + certificate, err := tls.LoadX509KeyPair(cfg.TLSCertPath, cfg.TLSKeyPath)
  85. + if err != nil {
  86. + panic(err)
  87. + }
  88. + config := &tls.Config{
  89. + Certificates: []tls.Certificate{certificate},
  90. + }
  91. +
  92. + // Create TLS listener.
  93. + tlslL := tls.NewListener(tlsl, config)
  94. +
  95. + // Serve HTTP over TLS.
  96. + go grace.Serve(tlslL)
  97. + }
  98. +
  99. grace.SetKeepAlivesEnabled(false)
  100.  
  101. + if err := m.Serve(); !strings.Contains(err.Error(), "use of closed network connection") {
  102. + panic(err)
  103. + }
  104. +
  105. +/*
  106. if err := grace.ListenAndServe(); err != nil {
  107. if opErr, ok := err.(*net.OpError); !ok || (ok && opErr.Op != "accept") {
  108. glog.Errorf("Failed to gracefully run HTTP server: %s", err.Error())
  109. }
  110. }
  111. + */
  112.  
  113. if err := srv.tracker.Close(); err != nil {
  114. glog.Errorf("Failed to shutdown tracker cleanly: %s", err.Error())
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement