From 366bddd20a40f7090c70fc7a4001e181a0145510 Mon Sep 17 00:00:00 2001
From: Joseph Schorr <joseph.schorr@coreos.com>
Date: Tue, 30 Oct 2018 15:39:57 -0400
Subject: [PATCH] Fix NPE bug in torrentinfo set call

Fixes https://sentry.io/coreos/backend-production/issues/747494304
---
 data/model/storage.py                          | 12 ++++++++----
 data/registry_model/test/test_pre_oci_model.py |  4 ++++
 2 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/data/model/storage.py b/data/model/storage.py
index 85a118e88..79cff50e1 100644
--- a/data/model/storage.py
+++ b/data/model/storage.py
@@ -338,10 +338,14 @@ def get_storage_locations(uuid):
 
 def save_torrent_info(storage_object, piece_length, pieces):
   try:
-    return TorrentInfo.create(storage=storage_object, piece_length=piece_length, pieces=pieces)
-  except IntegrityError:
-    # TorrentInfo already exists for this storage.
-    pass
+    return TorrentInfo.get(storage=storage_object, piece_length=piece_length)
+  except TorrentInfo.DoesNotExist:
+    try:
+      return TorrentInfo.create(storage=storage_object, piece_length=piece_length, pieces=pieces)
+    except IntegrityError:
+      # TorrentInfo already exists for this storage.
+      return TorrentInfo.get(storage=storage_object, piece_length=piece_length)
+
 
 def get_torrent_info(blob):
   try:
diff --git a/data/registry_model/test/test_pre_oci_model.py b/data/registry_model/test/test_pre_oci_model.py
index 7fa41b63f..addadf6cc 100644
--- a/data/registry_model/test/test_pre_oci_model.py
+++ b/data/registry_model/test/test_pre_oci_model.py
@@ -564,6 +564,10 @@ def test_torrent_info(pre_oci_model):
   assert pre_oci_model.get_torrent_info(layers[0].blob) is None
   pre_oci_model.set_torrent_info(layers[0].blob, 2, 'foo')
 
+  # Set it again exactly, which should be a no-op.
+  pre_oci_model.set_torrent_info(layers[0].blob, 2, 'foo')
+
+  # Check the information we've set.
   torrent_info = pre_oci_model.get_torrent_info(layers[0].blob)
   assert torrent_info is not None
   assert torrent_info.piece_length == 2