From 95797a82525a9e35a654a5375b42415672c07ffd Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Fri, 25 Mar 2022 14:24:10 +0200
Subject: [PATCH] Move reading of container state for internal/oci package

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 .../hook/update-ldcache/update-ldcache.go     | 27 +------
 internal/oci/state.go                         | 73 +++++++++++++++++++
 2 files changed, 76 insertions(+), 24 deletions(-)
 create mode 100644 internal/oci/state.go

diff --git a/cmd/nvidia-ctk/hook/update-ldcache/update-ldcache.go b/cmd/nvidia-ctk/hook/update-ldcache/update-ldcache.go
index 56f9f229..176cd777 100644
--- a/cmd/nvidia-ctk/hook/update-ldcache/update-ldcache.go
+++ b/cmd/nvidia-ctk/hook/update-ldcache/update-ldcache.go
@@ -17,14 +17,12 @@
 package ldcache
 
 import (
-	"encoding/json"
 	"fmt"
 	"os"
 	"path/filepath"
 	"syscall"
 
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/oci"
-	"github.com/opencontainers/runtime-spec/specs-go"
 	"github.com/sirupsen/logrus"
 	"github.com/urfave/cli/v2"
 )
@@ -76,31 +74,12 @@ func (m command) build() *cli.Command {
 }
 
 func (m command) run(c *cli.Context, cfg *config) error {
-	var s specs.State
-
-	inputReader := os.Stdin
-	if cfg.containerSpec != "" && cfg.containerSpec != "-" {
-		inputFile, err := os.Open(cfg.containerSpec)
-		if err != nil {
-			return fmt.Errorf("failed to open intput: %v", err)
-		}
-		defer inputFile.Close()
-		inputReader = inputFile
-	}
-
-	d := json.NewDecoder(inputReader)
-	if err := d.Decode(&s); err != nil {
-		return fmt.Errorf("failed to decode container state: %v", err)
-	}
-
-	specFilePath := oci.GetSpecFilePath(s.Bundle)
-	specFile, err := os.Open(specFilePath)
+	s, err := oci.LoadContainerState(cfg.containerSpec)
 	if err != nil {
-		return fmt.Errorf("failed to open OCI spec file: %v", err)
+		return fmt.Errorf("failed to load container state: %v", err)
 	}
-	defer specFile.Close()
 
-	spec, err := oci.LoadFrom(specFile)
+	spec, err := s.LoadSpec()
 	if err != nil {
 		return fmt.Errorf("failed to load OCI spec: %v", err)
 	}
diff --git a/internal/oci/state.go b/internal/oci/state.go
new file mode 100644
index 00000000..d3939b68
--- /dev/null
+++ b/internal/oci/state.go
@@ -0,0 +1,73 @@
+/**
+# Copyright (c) 2022, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+**/
+
+package oci
+
+import (
+	"encoding/json"
+	"fmt"
+	"io"
+	"os"
+
+	"github.com/opencontainers/runtime-spec/specs-go"
+)
+
+// State stores an OCI container state. This includes the spec path and the environment
+type State specs.State
+
+// LoadContainerState loads the container state from the specified filename. If the filename is empty or '-' the state is loaded from STDIN
+func LoadContainerState(filename string) (*State, error) {
+	if filename == "" || filename == "-" {
+		return ReadContainerState(os.Stdin)
+	}
+
+	inputFile, err := os.Open(filename)
+	if err != nil {
+		return nil, fmt.Errorf("failed to open file: %v", err)
+	}
+	defer inputFile.Close()
+
+	return ReadContainerState(inputFile)
+}
+
+// ReadContainerState reads the container state from the specified reader
+func ReadContainerState(reader io.Reader) (*State, error) {
+	var s State
+
+	d := json.NewDecoder(reader)
+	if err := d.Decode(&s); err != nil {
+		return nil, fmt.Errorf("failed to decode container state: %v", err)
+	}
+
+	return &s, nil
+}
+
+// LoadSpec loads the OCI spec associated with the container state
+func (s State) LoadSpec() (*specs.Spec, error) {
+	specFilePath := GetSpecFilePath(s.Bundle)
+	specFile, err := os.Open(specFilePath)
+	if err != nil {
+		return nil, fmt.Errorf("failed to open OCI spec file: %v", err)
+	}
+	defer specFile.Close()
+
+	spec, err := LoadFrom(specFile)
+	if err != nil {
+		return nil, fmt.Errorf("failed to load OCI spec: %v", err)
+	}
+
+	return spec, nil
+}