From fc581c378479c8d78daf9583a3a970482e7e4a3e Mon Sep 17 00:00:00 2001
From: Marwan Sulaiman <marwan.sameer@gmail.com>
Date: Fri, 24 Aug 2018 02:13:09 -0400
Subject: [PATCH] pkg/storage: ensure mongo returns KindNotFound (#558)

---
 pkg/storage/mongo/getter.go | 22 ++++++++++++++--------
 pkg/storage/mongo/lister.go |  9 +++++----
 2 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/pkg/storage/mongo/getter.go b/pkg/storage/mongo/getter.go
index 4ff038de..9b856ca0 100644
--- a/pkg/storage/mongo/getter.go
+++ b/pkg/storage/mongo/getter.go
@@ -3,8 +3,8 @@ package mongo
 import (
 	"context"
 	"io"
-	"strings"
 
+	"github.com/globalsign/mgo"
 	"github.com/globalsign/mgo/bson"
 	"github.com/gomods/athens/pkg/errors"
 	"github.com/gomods/athens/pkg/storage"
@@ -20,10 +20,11 @@ func (s *ModuleStore) Info(ctx context.Context, module, vsn string) ([]byte, err
 	result := &storage.Module{}
 	err := c.Find(bson.M{"module": module, "version": vsn}).One(result)
 	if err != nil {
-		if strings.Contains(err.Error(), "not found") {
-			err = errors.E(op, errors.M(module), errors.V(vsn), errors.KindNotFound)
+		kind := errors.KindUnexpected
+		if err == mgo.ErrNotFound {
+			kind = errors.KindNotFound
 		}
-		return nil, err
+		return nil, errors.E(op, kind, errors.M(module), errors.V(vsn), err)
 	}
 
 	return result.Info, nil
@@ -38,10 +39,11 @@ func (s *ModuleStore) GoMod(ctx context.Context, module, vsn string) ([]byte, er
 	result := &storage.Module{}
 	err := c.Find(bson.M{"module": module, "version": vsn}).One(result)
 	if err != nil {
-		if strings.Contains(err.Error(), "not found") {
-			err = errors.E(op, errors.M(module), errors.V(vsn), errors.KindNotFound)
+		kind := errors.KindUnexpected
+		if err == mgo.ErrNotFound {
+			kind = errors.KindNotFound
 		}
-		return nil, err
+		return nil, errors.E(op, kind, errors.M(module), errors.V(vsn), err)
 	}
 
 	return result.Mod, nil
@@ -57,7 +59,11 @@ func (s *ModuleStore) Zip(ctx context.Context, module, vsn string) (io.ReadClose
 	fs := s.s.DB(s.d).GridFS("fs")
 	f, err := fs.Open(zipName)
 	if err != nil {
-		return nil, errors.E(op, err)
+		kind := errors.KindUnexpected
+		if err == mgo.ErrNotFound {
+			kind = errors.KindNotFound
+		}
+		return nil, errors.E(op, err, kind)
 	}
 
 	return f, nil
diff --git a/pkg/storage/mongo/lister.go b/pkg/storage/mongo/lister.go
index 40aa14b4..9a4cb711 100644
--- a/pkg/storage/mongo/lister.go
+++ b/pkg/storage/mongo/lister.go
@@ -2,8 +2,8 @@ package mongo
 
 import (
 	"context"
-	"strings"
 
+	"github.com/globalsign/mgo"
 	"github.com/globalsign/mgo/bson"
 	"github.com/gomods/athens/pkg/errors"
 	"github.com/gomods/athens/pkg/storage"
@@ -19,10 +19,11 @@ func (s *ModuleStore) List(ctx context.Context, module string) ([]string, error)
 	result := make([]storage.Module, 0)
 	err := c.Find(bson.M{"module": module}).All(&result)
 	if err != nil {
-		if strings.Contains(err.Error(), "not found") {
-			err = errors.E(op, err, errors.M(module), errors.KindNotFound)
+		kind := errors.KindUnexpected
+		if err == mgo.ErrNotFound {
+			kind = errors.KindNotFound
 		}
-		return nil, err
+		return nil, errors.E(op, kind, errors.M(module), err)
 	}
 
 	versions := make([]string, len(result))