package whohas import ( "context" "mime" "net/http" "path/filepath" "runtime" "strconv" "strings" "sync" "time" ) type BACKEND struct { Host string Secure bool Path string Headers map[string]string Penalty time.Duration } type CACHE struct { TTL time.Duration last time.Time items map[string]*LOOKUP sync.RWMutex } type LOOKUP struct { index int deadline time.Time Protocol string Host string Headers map[string]string Size int64 Mime string Ranges bool Date time.Time Modified time.Time Expires time.Time } var cores int func Lookup(path string, backends []BACKEND, timeout time.Duration, cache *CACHE, ckey string) (lookup *LOOKUP) { if cores == 0 { cores = runtime.NumCPU() } if path == "" || backends == nil || len(backends) < 1 || timeout < 100*time.Millisecond { return } cpath := path if index := strings.Index(path, "?"); index >= 0 { cpath = path[:index] } cbackends := backends if cache != nil && cache.items != nil { now := time.Now() if cores > 1 { cache.RLock() } if cache.items[cpath] != nil && now.Sub(cache.items[cpath].deadline) < 0 { lookup = cache.items[cpath] if cache.items[cpath].Host != "" { cbackends = []BACKEND{} for _, backend := range backends { if backend.Host == cache.items[cpath].Host { cbackends = append(cbackends, backend) break } } if len(cbackends) < 1 { cbackends = backends } } } if len(cbackends) == len(backends) && ckey != "" && cache.items["k"+ckey] != nil && now.Sub(cache.items["k"+ckey].deadline) < 0 { cbackends = []BACKEND{} for _, backend := range backends { if backend.Host == cache.items["k"+ckey].Host { cbackends = append(cbackends, backend) break } } if len(cbackends) < 1 { cbackends = backends } } if cores > 1 { cache.RUnlock() } } if lookup == nil { inflight := len(cbackends) sink := make(chan LOOKUP, inflight+1) cancels := make([]context.CancelFunc, inflight) for index, backend := range cbackends { var ctx context.Context ctx, cancels[index] = context.WithCancel(context.Background()) go func(index int, backend BACKEND, ctx context.Context) { lookup := LOOKUP{index: index} if backend.Penalty != 0 && len(cbackends) > 1 { select { case <-time.After(backend.Penalty): case <-ctx.Done(): sink <- lookup return } } lookup.Protocol = "http" if backend.Secure { lookup.Protocol = "https" } rpath := path if backend.Path != "" { rpath = backend.Path } if request, err := http.NewRequest(http.MethodHead, lookup.Protocol+"://"+backend.Host+rpath, nil); err == nil { request = request.WithContext(ctx) request.Header.Set("User-Agent", "whohas") if backend.Headers != nil { lookup.Headers = map[string]string{} for name, value := range backend.Headers { lookup.Headers[name] = value request.Header.Set(name, value) } } if response, err := http.DefaultClient.Do(request); err == nil { if response.StatusCode == 200 { lookup.Host = backend.Host lookup.Size, _ = strconv.ParseInt(response.Header.Get("Content-Length"), 10, 64) lookup.Mime = response.Header.Get("Content-Type") if lookup.Mime == "" || lookup.Mime == "application/octet-stream" || lookup.Mime == "text/plain" { if extension := filepath.Ext(path); extension != "" { lookup.Mime = mime.TypeByExtension(extension) } } if response.Header.Get("Accept-Ranges") != "" { lookup.Ranges = true } if header := response.Header.Get("Date"); header != "" { lookup.Date, _ = http.ParseTime(header) } else { lookup.Date = time.Now() } if header := response.Header.Get("Last-Modified"); header != "" { lookup.Modified, _ = http.ParseTime(header) } if header := response.Header.Get("Expires"); header != "" { lookup.Expires, _ = http.ParseTime(header) } else { lookup.Expires = lookup.Date.Add(time.Hour) } if lookup.Expires.Sub(lookup.Date) < 2*time.Second { lookup.Expires = lookup.Date.Add(2 * time.Second) } } response.Body.Close() } } sink <- lookup }(index, backend, ctx) } for inflight > 0 { select { case result := <-sink: inflight-- cancels[result.index] = nil if result.Host != "" { lookup = &result for index, cancel := range cancels { if cancels[index] != nil && index != result.index { cancel() cancels[index] = nil } } } case <-time.After(timeout): for index, cancel := range cancels { if cancels[index] != nil { cancel() } } } } close(sink) } if cache != nil { now := time.Now() if cores > 1 { cache.Lock() } if cache.items == nil { cache.items = map[string]*LOOKUP{} } if now.Sub(cache.last) >= 5*time.Second { cache.last = now for key, item := range cache.items { if now.Sub(item.deadline) >= 0 { delete(cache.items, key) } } } if lookup == nil || lookup.Host == "" { if ckey != "" { delete(cache.items, "k"+ckey) } if cache.items[cpath] == nil { cache.items[cpath] = &LOOKUP{deadline: now.Add(5 * time.Second)} } lookup = nil } else { if cache.TTL < 2*time.Second { cache.TTL = 2 * time.Second } if ckey != "" { cache.items["k"+ckey] = &LOOKUP{Host: lookup.Host, deadline: now.Add(cache.TTL)} } if cache.items[cpath] == nil { ttl := lookup.Expires.Sub(lookup.Date) if ttl < 2*time.Second { ttl = 2 * time.Second } if ttl > cache.TTL { ttl = cache.TTL } if ttl > 10*time.Minute { ttl = 10 * time.Minute } lookup.deadline = now.Add(ttl) cache.items[cpath] = lookup } } if cores > 1 { cache.Unlock() } } return }