From f7c4aa78652fd7211c9b4f41ddc897cd956e43e0 Mon Sep 17 00:00:00 2001
From: Paul Duncan <pabs@pablotron.org>
Date: Tue, 22 Feb 2022 21:58:02 -0500
Subject: nvdmirror: use atomictemp

---
 nvdmirror/jsoncache.go | 35 +++++++++++++++++------------------
 nvdmirror/nvdmirror.go | 25 +++++++++----------------
 2 files changed, 26 insertions(+), 34 deletions(-)

diff --git a/nvdmirror/jsoncache.go b/nvdmirror/jsoncache.go
index 7e84a94..7bcd0ca 100644
--- a/nvdmirror/jsoncache.go
+++ b/nvdmirror/jsoncache.go
@@ -3,9 +3,10 @@ package nvdmirror
 import (
   "compress/gzip"
   "encoding/json"
+  "io"
   "os"
-  "path/filepath"
   "time"
+  "github.com/pablotron/cvez/atomictemp"
 )
 
 // JSON file backed URL cache.
@@ -102,26 +103,24 @@ func (me *JsonCache) Close() error {
   }
 
   // open temp output file
-  t, err := os.CreateTemp(filepath.Dir(me.path), "")
-  if err != nil {
-    return err
-  }
-  defer os.Remove(t.Name())
+  err := atomictemp.Create(me.path, func(w io.Writer) error {
+    // encode/compress JSON
+    zw := gzip.NewWriter(w)
+    e := json.NewEncoder(zw)
+    if err := e.Encode(me.vals); err != nil {
+      return err
+    }
 
-  // encode/compress JSON
-  zw := gzip.NewWriter(t)
-  e := json.NewEncoder(zw)
-  if err = e.Encode(me.vals); err != nil {
-    return err
-  }
+    // flush changes
+    if err := zw.Flush(); err != nil {
+      return err
+    }
 
-  // flush changes
-  if err = zw.Flush(); err != nil {
-    return err
-  }
+    // return success
+    return nil
+  })
 
-  // rename to destination file
-  if err = os.Rename(t.Name(), me.path); err != nil {
+  if err != nil {
     return err
   }
 
diff --git a/nvdmirror/nvdmirror.go b/nvdmirror/nvdmirror.go
index da758ad..cba638b 100644
--- a/nvdmirror/nvdmirror.go
+++ b/nvdmirror/nvdmirror.go
@@ -3,11 +3,11 @@ package nvdmirror
 
 import (
   "fmt"
+  "github.com/pablotron/cvez/atomictemp"
   "github.com/rs/zerolog/log"
   "io"
   "net/http"
   "net/url"
-  "os"
   "path/filepath"
   "time"
 )
@@ -83,14 +83,6 @@ func fetch(ch chan fetchResult, config SyncConfig, cache Cache, client *http.Cli
   path := filepath.Join(dstDir, filepath.Base(src.Path))
   log.Debug().Str("url", srcUrl).Str("path", path).Send()
 
-  // create temporary output file
-  f, err := os.CreateTemp(filepath.Dir(path), "")
-  if err != nil {
-    ch <- fetchResult { src: srcUrl, err: err }
-    return
-  }
-  defer os.Remove(f.Name())
-
   // create request
   req, err := http.NewRequest("GET", srcUrl, nil)
   if err != nil {
@@ -116,15 +108,16 @@ func fetch(ch chan fetchResult, config SyncConfig, cache Cache, client *http.Cli
 
   switch resp.StatusCode {
   case 200:
-    // copy body to result
-    if size, err := io.Copy(f, resp.Body); err != nil {
-      // copy failed
-      ch <- fetchResult { src: srcUrl, err: err, modified: false }
-    } else if err = os.Rename(f.Name(), path); err != nil {
-      // rename failed
+    // create temporary output file
+    err := atomictemp.Create(path, func(f io.Writer) error {
+      _, err := io.Copy(f, resp.Body)
+      return err
+    })
+
+    if err != nil {
+      // write failed
       ch <- fetchResult { src: srcUrl, err: err, modified: false }
     } else {
-      log.Debug().Str("url", srcUrl).Int64("size", size).Send()
       ch <- fetchResult {
         src: srcUrl,
         modified: true,
-- 
cgit v1.2.3