storage.go 1.4 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677
  1. package hpkp
  2. import (
  3. "strings"
  4. "sync"
  5. )
  6. // MemStorage is threadsafe hpkp host storage backed by an in-memory map
  7. type MemStorage struct {
  8. domains map[string]Header
  9. mutex sync.Mutex
  10. }
  11. // NewMemStorage initializes hpkp in-memory datastructure
  12. func NewMemStorage() *MemStorage {
  13. m := &MemStorage{}
  14. m.domains = make(map[string]Header)
  15. return m
  16. }
  17. // Lookup returns the corresponding hpkp header information for a given host
  18. func (s *MemStorage) Lookup(host string) *Header {
  19. s.mutex.Lock()
  20. defer s.mutex.Unlock()
  21. d, ok := s.domains[host]
  22. if ok {
  23. return copy(d)
  24. }
  25. // is h a subdomain of an hpkp domain, walk the domain to see if it is a sub
  26. // sub ... sub domain of a domain that has the `includeSubDomains` rule
  27. l := len(host)
  28. for l > 0 {
  29. i := strings.Index(host, ".")
  30. if i > 0 {
  31. host = host[i+1:]
  32. d, ok := s.domains[host]
  33. if ok {
  34. if d.IncludeSubDomains {
  35. return copy(d)
  36. }
  37. }
  38. l = len(host)
  39. } else {
  40. break
  41. }
  42. }
  43. return nil
  44. }
  45. func copy(h Header) *Header {
  46. d := h
  47. return &d
  48. }
  49. // Add a domain to hpkp storage
  50. func (s *MemStorage) Add(host string, d *Header) {
  51. s.mutex.Lock()
  52. defer s.mutex.Unlock()
  53. if s.domains == nil {
  54. s.domains = make(map[string]Header)
  55. }
  56. if d.MaxAge == 0 && !d.Permanent {
  57. check, ok := s.domains[host]
  58. if ok {
  59. if !check.Permanent {
  60. delete(s.domains, host)
  61. }
  62. }
  63. } else {
  64. s.domains[host] = *d
  65. }
  66. }