From c3a29cce8b05eb89159e4bfacdb579aa6edc15eb Mon Sep 17 00:00:00 2001
From: Tonis Tiigi <tonistiigi@gmail.com>
Date: Mon, 9 Jul 2018 14:26:31 -0700
Subject: [PATCH] gateway: implementation for return

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>
---
 frontend/gateway/gateway.go | 46 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 45 insertions(+), 1 deletion(-)

diff --git a/frontend/gateway/gateway.go b/frontend/gateway/gateway.go
index 083677c1..9f507792 100644
--- a/frontend/gateway/gateway.go
+++ b/frontend/gateway/gateway.go
@@ -194,6 +194,13 @@ func (gf *gatewayFrontend) Solve(ctx context.Context, llbBridge frontend.Fronten
 
 	defer func() {
 		for _, r := range lbf.refs {
+			if retErr != nil && len(lbf.lastRefs) > 0 {
+				for _, ref := range lbf.lastRefs {
+					if ref == r {
+						continue
+					}
+				}
+			}
 			if r != nil && (lbf.lastRef != r || retErr != nil) {
 				r.Release(context.TODO())
 			}
@@ -209,10 +216,18 @@ func (gf *gatewayFrontend) Solve(ctx context.Context, llbBridge frontend.Fronten
 		ReadonlyRootFS: readonly,
 	}, rootFS, lbf.Stdin, lbf.Stdout, os.Stderr)
 
+	if lbf.err != nil {
+		return nil, nil, lbf.err
+	}
+
 	if err != nil {
 		return nil, nil, err
 	}
 
+	if lbf.lastRefs != nil {
+		return lbf.lastRefs, lbf.exporterAttr, nil
+	}
+
 	return map[string]solver.CachedResult{"default": lbf.lastRef}, lbf.exporterAttr, nil
 }
 
@@ -293,6 +308,8 @@ type llbBridgeForwarder struct {
 	llbBridge    frontend.FrontendLLBBridge
 	refs         map[string]solver.CachedResult
 	lastRef      solver.CachedResult
+	lastRefs     map[string]solver.CachedResult
+	err          error
 	exporterAttr map[string][]byte
 	workers      frontend.WorkerInfos
 	*pipe
@@ -432,7 +449,34 @@ func (lbf *llbBridgeForwarder) Ping(context.Context, *pb.PingRequest) (*pb.PongR
 }
 
 func (lbf *llbBridgeForwarder) Return(ctx context.Context, in *pb.ReturnRequest) (*pb.ReturnResponse, error) {
-	return nil, errors.Errorf("return not implemented yet")
+	if in.Error != nil {
+		lbf.err = errors.Errorf(in.Error.Text)
+	} else {
+		lbf.exporterAttr = in.ExporterMeta
+		refs, err := lbf.convertRefs(in.ExporterRefs)
+		if err != nil {
+			lbf.err = err
+		} else {
+			lbf.lastRefs = refs
+		}
+	}
+
+	return &pb.ReturnResponse{}, nil
+}
+
+func (lbf *llbBridgeForwarder) convertRefs(refs map[string]string) (map[string]solver.CachedResult, error) {
+	out := make(map[string]solver.CachedResult, len(refs))
+	for k, v := range refs {
+		if v == "" {
+			out[k] = nil
+		}
+		r, ok := lbf.refs[v]
+		if !ok {
+			return nil, errors.Errorf("return reference %s not found", v)
+		}
+		out[k] = r
+	}
+	return out, nil
 }
 
 func serve(ctx context.Context, grpcServer *grpc.Server, conn net.Conn) {