middleware.go 2.0KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586
  1. package livereload
  2. import (
  3. "bytes"
  4. "compress/gzip"
  5. "compress/zlib"
  6. "fmt"
  7. "io"
  8. "net/http"
  9. "regexp"
  10. "strings"
  11. "github.com/goanywhere/env"
  12. )
  13. type writer struct {
  14. http.ResponseWriter
  15. host string
  16. }
  17. func (self *writer) addJavaScript(data []byte) []byte {
  18. javascript := fmt.Sprintf(`<script src="//%s%s"></script>
  19. </head>`, self.host, URL.JavaScript)
  20. return regexp.MustCompile(`</head>`).ReplaceAll(data, []byte(javascript))
  21. }
  22. func (self *writer) Write(data []byte) (size int, e error) {
  23. if strings.Contains(self.Header().Get("Content-Type"), "html") {
  24. var encoding = self.Header().Get("Content-Encoding")
  25. if encoding == "" {
  26. data = self.addJavaScript(data)
  27. } else {
  28. var reader io.ReadCloser
  29. var buffer *bytes.Buffer = new(bytes.Buffer)
  30. if encoding == "gzip" {
  31. // decode to add javascript reference.
  32. reader, _ = gzip.NewReader(bytes.NewReader(data))
  33. io.Copy(buffer, reader)
  34. output := self.addJavaScript(buffer.Bytes())
  35. reader.Close()
  36. buffer.Reset()
  37. // encode back to HTML with added javascript reference.
  38. writer := gzip.NewWriter(buffer)
  39. writer.Write(output)
  40. writer.Close()
  41. data = buffer.Bytes()
  42. } else if encoding == "deflate" {
  43. // decode to add javascript reference.
  44. reader, _ = zlib.NewReader(bytes.NewReader(data))
  45. io.Copy(buffer, reader)
  46. output := self.addJavaScript(buffer.Bytes())
  47. reader.Close()
  48. buffer.Reset()
  49. // encode back to HTML with added javascript reference.
  50. writer := zlib.NewWriter(buffer)
  51. writer.Write(output)
  52. writer.Close()
  53. data = buffer.Bytes()
  54. }
  55. }
  56. }
  57. return self.ResponseWriter.Write(data)
  58. }
  59. func Module(next http.Handler) http.Handler {
  60. // ONLY run this under debug mode.
  61. if !env.Bool("DEBUG", true) {
  62. return next
  63. }
  64. Start()
  65. fn := func(w http.ResponseWriter, r *http.Request) {
  66. if r.URL.Path == URL.WebSocket {
  67. ServeWebSocket(w, r)
  68. } else if r.URL.Path == URL.JavaScript {
  69. ServeJavaScript(w, r)
  70. } else {
  71. writer := &writer{w, r.Host}
  72. next.ServeHTTP(writer, r)
  73. }
  74. }
  75. return http.HandlerFunc(fn)
  76. }