nvdmirror: add SyncConfig, add additional tests

2 files changed, 123 insertions, 55 deletions

diff --git a/nvdmirror/nvdmirror.go b/nvdmirror/nvdmirror.go
index 7b092b1..da758ad 100644
--- a/nvdmirror/nvdmirror.go
+++ b/nvdmirror/nvdmirror.go
@@ -12,28 +12,46 @@ import (
   "time"
 )
 
-// URLs
-type Urls struct {
-  // CVE 1.1 Base URL.  The full meta and JSON URLs are constructed from
-  // this appended to this base..
-  Cve11Base string
+// default user agent (FIXME: make configurable)
+var defaultUserAgent = "cvez/0.1.0"
+
+// Sync() configuration.
+type SyncConfig struct {
+  // CVE 1.1 Base URL.  The full meta and JSON URLs are constructed by
+  // appending the file name to this base.
+  Cve11BaseUrl string
+
+  // CPE 1.0 base URL.  The full meta and JSON URLs are constructed by
+  // appending the file name to this base.
+  Cpe10MatchBaseUrl string
 
-  // CPE 1.0 match meta URL
-  Cpe10MatchMeta string
+  // CPE 2.3 dictionary URL.
+  Cpe23DictUrl string
 
-  // CPE 1.0 match URL
-  Cpe10MatchJson string
+  // User agent string.  Set to "" for default user agent string.
+  UserAgent string
 
-  // CPE 2.3 dictionary URL
-  Cpe23Dict string
+  // Maximum number of idle connections.
+	MaxIdleConns int
+
+  // Idle connection timeout.
+	IdleConnTimeout time.Duration
+}
+
+// Get user agent string.
+func (me SyncConfig) GetUserAgent() string {
+  if len(me.UserAgent) > 0 {
+    return me.UserAgent
+  } else {
+    return defaultUserAgent
+  }
 }
 
 // NVD URLs
-var DefaultUrls = Urls {
-  Cve11Base: "https://nvd.nist.gov/feeds/json/cve/1.1",
-  Cpe10MatchMeta: "https://nvd.nist.gov/feeds/json/cpematch/1.0/nvdcpematch-1.0.meta",
-  Cpe10MatchJson: "https://nvd.nist.gov/feeds/json/cpematch/1.0/nvdcpematch-1.0.gz",
-  Cpe23Dict: "https://nvd.nist.gov/feeds/xml/cpe/dictionary/official-cpe-dictionary_v2.3.xml.gz",
+var DefaultConfig = SyncConfig {
+  Cve11BaseUrl: "https://nvd.nist.gov/feeds/json/cve/1.1",
+  Cpe10MatchBaseUrl: "https://nvd.nist.gov/feeds/json/cpematch/1.0",
+  Cpe23DictUrl: "https://nvd.nist.gov/feeds/xml/cpe/dictionary/official-cpe-dictionary_v2.3.xml.gz",
 }
 
 var extraFiles = []string {
@@ -45,7 +63,7 @@ var extraFiles = []string {
 type fetchResult struct {
   src string // source URL
   err error // fetch result
-  notModified bool // was the result unmodified?
+  modified bool // was the result modified?
   lastModified string // last modified response header
   etag string // etag response header
 }
@@ -53,10 +71,7 @@ type fetchResult struct {
 // base CVE year
 const baseYear = 2002
 
-// default user agent (FIXME: make configurable)
-var userAgent = "cvez/0.1.0"
-
-func fetch(ch chan fetchResult, client *http.Client, cache Cache, srcUrl string, dstDir string) {
+func fetch(ch chan fetchResult, config SyncConfig, cache Cache, client *http.Client, srcUrl string, dstDir string) {
   // parse source url
   src, err := url.Parse(srcUrl)
   if err != nil {
@@ -84,7 +99,7 @@ func fetch(ch chan fetchResult, client *http.Client, cache Cache, srcUrl string,
   }
 
   // add request headers
-  req.Header.Add("user-agent", userAgent)
+  req.Header.Add("user-agent", config.GetUserAgent())
   if headers, ok := cache.Get(srcUrl); ok {
     for k, v := range(headers) {
       req.Header.Add(k, v)
@@ -104,20 +119,21 @@ func fetch(ch chan fetchResult, client *http.Client, cache Cache, srcUrl string,
     // copy body to result
     if size, err := io.Copy(f, resp.Body); err != nil {
       // copy failed
-      ch <- fetchResult { src: srcUrl, err: err }
+      ch <- fetchResult { src: srcUrl, err: err, modified: false }
     } else if err = os.Rename(f.Name(), path); err != nil {
       // rename failed
-      ch <- fetchResult { src: srcUrl, err: err }
+      ch <- fetchResult { src: srcUrl, err: err, modified: false }
     } else {
       log.Debug().Str("url", srcUrl).Int64("size", size).Send()
       ch <- fetchResult {
         src: srcUrl,
+        modified: true,
         lastModified: resp.Header.Get("last-modified"),
         etag: resp.Header.Get("etag"),
       }
     }
   case 304:
-    ch <- fetchResult { src: srcUrl, notModified: true }
+    ch <- fetchResult { src: srcUrl, modified: false }
   default:
     code := resp.StatusCode
     err := fmt.Errorf("%d: %s", code, http.StatusText(code))
@@ -125,7 +141,7 @@ func fetch(ch chan fetchResult, client *http.Client, cache Cache, srcUrl string,
   }
 }
 
-func Sync(urls Urls, cache Cache, dstDir string) error {
+func Sync(config SyncConfig, cache Cache, dstDir string) error {
   log.Debug().Str("dstDir", dstDir).Msg("Sync")
 
   // create fetch result channel
@@ -133,27 +149,27 @@ func Sync(urls Urls, cache Cache, dstDir string) error {
 
   // create shared transport and client
   tr := &http.Transport{
-    // FIXME: make configurable
-	  MaxIdleConns:     10,
-	  IdleConnTimeout:  30 * time.Second,
+	  MaxIdleConns:     config.MaxIdleConns,
+	  IdleConnTimeout:  config.IdleConnTimeout,
   }
   client := &http.Client{Transport: tr}
 
   // calculate total number of years
-  numYears := time.Now().Year() - 2002 + 1
+  numYears := time.Now().Year() - baseYear + 1
 
-  // fetch metas
+  // fetch cve feed metas
   for year := baseYear; year < baseYear + numYears; year++ {
     // build url
-    url := fmt.Sprintf("%s/nvdcve-1.1-%04d.meta", urls.Cve11Base, year)
+    url := fmt.Sprintf("%s/nvdcve-1.1-%04d.meta", config.Cve11BaseUrl, year)
     log.Debug().Int("year", year).Str("url", url).Send()
-    go fetch(ch, client, cache, url, dstDir)
+    go fetch(ch, config, cache, client, url, dstDir)
   }
 
+  // fetch cve extra file metas
   for _, s := range(extraFiles) {
-    url := fmt.Sprintf("%s/%s.meta", urls.Cve11Base, s)
+    url := fmt.Sprintf("%s/%s.meta", config.Cve11BaseUrl, s)
     log.Debug().Str("file", s).Str("url", url).Send()
-    go fetch(ch, client, cache, url, dstDir)
+    go fetch(ch, config, cache, client, url, dstDir)
   }
 
   // read results
@@ -164,14 +180,15 @@ func Sync(urls Urls, cache Cache, dstDir string) error {
     } else {
       log.Info().Str("url", r.src).Str("etag", r.etag).Msg("ok")
 
-      if !r.notModified {
-        // cache headers
-        err := cache.Set(r.src, map[string]string {
+      if r.modified {
+        // build request headers
+        headers := map[string]string {
           "if-none-match": r.etag,
           "if-modified-since": r.lastModified,
-        })
+        }
 
-        if err != nil {
+        // save headers to cache
+        if err := cache.Set(r.src, headers); err != nil {
           log.Error().Str("url", r.src).Err(r.err).Msg("Set")
           return err
         }
@@ -179,6 +196,10 @@ func Sync(urls Urls, cache Cache, dstDir string) error {
     }
   }
 
+  // TODO: fetch cpe dictionary
+  // TODO: fetch cpematch
+  // "nvdcpematch-1.0.{meta,json}",
+
   // return success
   return nil
 }
diff --git a/nvdmirror/nvdmirror_test.go b/nvdmirror/nvdmirror_test.go
index 5b40dd8..2c65d41 100644
--- a/nvdmirror/nvdmirror_test.go
+++ b/nvdmirror/nvdmirror_test.go
@@ -60,22 +60,67 @@ func TestSync(t *testing.T) {
 
   // custom sync config
   // FIXME: stand up custom server for this
-  urls := Urls {
-    Cve11Base: fmt.Sprintf("http://localhost:%d", port),
+  config := SyncConfig {
+    Cve11BaseUrl: fmt.Sprintf("http://localhost:%d", port),
   }
 
   // sync data
-  if err := Sync(urls, &cache, dir); err != nil {
-    t.Error(err)
-  }
+  t.Run("initial", func(t *testing.T) {
+    if err := Sync(config, &cache, dir); err != nil {
+      t.Error(err)
+    }
+  })
 
   // sync data again (to test caching)
-  if err := Sync(urls, &cache, dir); err != nil {
-    t.Error(err)
-  }
+  t.Run("caching", func(t *testing.T) {
+    if err := Sync(config, &cache, dir); err != nil {
+      t.Error(err)
+    }
+  })
+
+  // sync w/ missing dir
+  t.Run("missingDir", func(t *testing.T) {
+    missingDir := filepath.Join(dir, "does/not/exist")
+    if err := Sync(config, &cache, missingDir); err != nil {
+      t.Error(err)
+    }
+  })
+
+  // sync w/ bad cache
+  t.Run("failSetCache", func(t *testing.T) {
+    var cache FailSetCache
+    if err := Sync(config, &cache, dir); err == nil {
+      t.Error(err)
+    }
+  })
+
+  t.Run("customUserAgent", func(t *testing.T) {
+    // custom sync config
+    // FIXME: stand up custom server for this
+    config := SyncConfig {
+      Cve11BaseUrl: fmt.Sprintf("http://localhost:%d", port),
+      UserAgent: "custom-user-agent/0.0.0",
+    }
+
+    if err := Sync(config, &cache, dir); err != nil {
+      t.Error(err)
+    }
+  })
+
+  t.Run("clientFail", func(t *testing.T) {
+    // custom sync config
+    // FIXME: stand up custom server for this
+    config := SyncConfig {
+      Cve11BaseUrl: "http://localhost:0",
+    }
+
+    if err := Sync(config, &cache, dir); err != nil {
+      t.Error(err)
+    }
+  })
 }
 
-func TestBadUrl(t *testing.T) {
+func TestBadUrls(t *testing.T) {
   // create temp dir
   dir, err := os.MkdirTemp("", "")
   if err != nil {
@@ -92,19 +137,21 @@ func TestBadUrl(t *testing.T) {
   }
   defer cache.Close()
 
+  // invalid base URLs
   failTests := []string {
-    "httsldkfjasldkfjadp://localhost:0",
-    "http://localhost:0",
+    "httsldkfjasldkfjadp:// \\ localhost:0",
   }
 
   for _, test := range(failTests) {
     t.Run(test, func(t *testing.T) {
       // custom sync config
-      urls := Urls { Cve11Base: test }
+      config := SyncConfig { Cve11BaseUrl: test }
 
-      // sync data
-      if err := Sync(urls, &cache, dir); err == nil {
-        t.Errorf("got success, exp error")
+      // sync data; note: even with an invalid base URL we still expect
+      // this call to succeed; it's just that all of the URLs will be
+      // nonsensical
+      if err := Sync(config, &cache, dir); err != nil {
+        t.Error(err)
       }
     })
   }