exec.go

Documentation: oss.terrastruct.com/d2/d2plugin

     1  package d2plugin
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"encoding/json"
     7  	"errors"
     8  	"fmt"
     9  	"os/exec"
    10  	"strconv"
    11  	"strings"
    12  	"time"
    13  
    14  	"oss.terrastruct.com/util-go/xdefer"
    15  	"oss.terrastruct.com/util-go/xmain"
    16  
    17  	"oss.terrastruct.com/d2/d2graph"
    18  	timelib "oss.terrastruct.com/d2/lib/time"
    19  )
    20  
    21  // execPlugin uses the binary at pathname with the plugin protocol to implement
    22  // the Plugin interface.
    23  //
    24  // The layout plugin protocol works as follows.
    25  //
    26  // Info
    27  //  1. The binary is invoked with info as the first argument.
    28  //  2. The stdout of the binary is unmarshalled into PluginInfo.
    29  //
    30  // Layout
    31  //  1. The binary is invoked with layout as the first argument and the json marshalled
    32  //     d2graph.Graph on stdin.
    33  //  2. The stdout of the binary is unmarshalled into a d2graph.Graph
    34  //
    35  // PostProcess
    36  //  1. The binary is invoked with postprocess as the first argument and the
    37  //     bytes of the SVG render on stdin.
    38  //  2. The stdout of the binary is bytes of SVG with any post-processing.
    39  //
    40  // If any errors occur the binary will exit with a non zero status code and write
    41  // the error to stderr.
    42  type execPlugin struct {
    43  	path string
    44  	opts map[string]string
    45  	info *PluginInfo
    46  }
    47  
    48  func (p *execPlugin) Flags(ctx context.Context) (_ []PluginSpecificFlag, err error) {
    49  	ctx, cancel := context.WithTimeout(ctx, time.Second*10)
    50  	defer cancel()
    51  	cmd := exec.CommandContext(ctx, p.path, "flags")
    52  	defer xdefer.Errorf(&err, "failed to run %v", cmd.Args)
    53  
    54  	stdout, err := cmd.Output()
    55  	if err != nil {
    56  		ee := &exec.ExitError{}
    57  		if errors.As(err, &ee) && len(ee.Stderr) > 0 {
    58  			return nil, fmt.Errorf("%v\nstderr:\n%s", ee, ee.Stderr)
    59  		}
    60  		return nil, err
    61  	}
    62  
    63  	var flags []PluginSpecificFlag
    64  
    65  	err = json.Unmarshal(stdout, &flags)
    66  	if err != nil {
    67  		return nil, fmt.Errorf("failed to unmarshal json: %w", err)
    68  	}
    69  
    70  	return flags, nil
    71  }
    72  
    73  func (p *execPlugin) HydrateOpts(opts []byte) error {
    74  	if opts != nil {
    75  		var execOpts map[string]interface{}
    76  		err := json.Unmarshal(opts, &execOpts)
    77  		if err != nil {
    78  			return xmain.UsageErrorf("non-exec layout options given for exec")
    79  		}
    80  
    81  		allString := make(map[string]string)
    82  		for k, v := range execOpts {
    83  			switch vt := v.(type) {
    84  			case string:
    85  				allString[k] = vt
    86  			case int64:
    87  				allString[k] = strconv.Itoa(int(vt))
    88  			case []interface{}:
    89  				str := make([]string, len(vt))
    90  				for i, v := range vt {
    91  					switch vt := v.(type) {
    92  					case string:
    93  						str[i] = vt
    94  					case int64:
    95  						str[i] = strconv.Itoa(int(vt))
    96  					case float64:
    97  						str[i] = strconv.Itoa(int(vt))
    98  					}
    99  				}
   100  				allString[k] = strings.Join(str, ",")
   101  			case []int64:
   102  				str := make([]string, len(vt))
   103  				for i, v := range vt {
   104  					str[i] = strconv.Itoa(int(v))
   105  				}
   106  				allString[k] = strings.Join(str, ",")
   107  			case float64:
   108  				allString[k] = strconv.Itoa(int(vt))
   109  			}
   110  		}
   111  
   112  		p.opts = allString
   113  	}
   114  	return nil
   115  }
   116  
   117  func (p *execPlugin) Info(ctx context.Context) (_ *PluginInfo, err error) {
   118  	if p.info != nil {
   119  		return p.info, nil
   120  	}
   121  
   122  	ctx, cancel := context.WithTimeout(ctx, time.Second*10)
   123  	defer cancel()
   124  	cmd := exec.CommandContext(ctx, p.path, "info")
   125  	defer xdefer.Errorf(&err, "failed to run %v", cmd.Args)
   126  
   127  	stdout, err := cmd.Output()
   128  	if err != nil {
   129  		ee := &exec.ExitError{}
   130  		if errors.As(err, &ee) && len(ee.Stderr) > 0 {
   131  			return nil, fmt.Errorf("%v\nstderr:\n%s", ee, ee.Stderr)
   132  		}
   133  		return nil, err
   134  	}
   135  
   136  	var info PluginInfo
   137  
   138  	err = json.Unmarshal(stdout, &info)
   139  	if err != nil {
   140  		return nil, fmt.Errorf("failed to unmarshal json: %w", err)
   141  	}
   142  
   143  	info.Type = "binary"
   144  	info.Path = p.path
   145  
   146  	p.info = &info
   147  	return &info, nil
   148  }
   149  
   150  func (p *execPlugin) Layout(ctx context.Context, g *d2graph.Graph) error {
   151  	ctx, cancel := timelib.WithTimeout(ctx, time.Minute*2)
   152  	defer cancel()
   153  
   154  	graphBytes, err := d2graph.SerializeGraph(g)
   155  	if err != nil {
   156  		return err
   157  	}
   158  
   159  	args := []string{"layout"}
   160  	for k, v := range p.opts {
   161  		args = append(args, fmt.Sprintf("--%s", k), v)
   162  	}
   163  	cmd := exec.CommandContext(ctx, p.path, args...)
   164  
   165  	buffer := bytes.Buffer{}
   166  	buffer.Write(graphBytes)
   167  	cmd.Stdin = &buffer
   168  
   169  	stdout, err := cmd.Output()
   170  	if err != nil {
   171  		ee := &exec.ExitError{}
   172  		if errors.As(err, &ee) && len(ee.Stderr) > 0 {
   173  			return fmt.Errorf("%v\nstderr:\n%s", ee, ee.Stderr)
   174  		}
   175  		return err
   176  	}
   177  	err = d2graph.DeserializeGraph(stdout, g)
   178  	if err != nil {
   179  		return fmt.Errorf("failed to unmarshal json: %w", err)
   180  	}
   181  
   182  	return nil
   183  }
   184  
   185  func (p *execPlugin) PostProcess(ctx context.Context, in []byte) ([]byte, error) {
   186  	ctx, cancel := context.WithTimeout(ctx, time.Minute)
   187  	defer cancel()
   188  
   189  	cmd := exec.CommandContext(ctx, p.path, "postprocess")
   190  
   191  	cmd.Stdin = bytes.NewBuffer(in)
   192  
   193  	stdout, err := cmd.Output()
   194  	if err != nil {
   195  		ee := &exec.ExitError{}
   196  		if errors.As(err, &ee) && len(ee.Stderr) > 0 {
   197  			return nil, fmt.Errorf("%v\nstderr:\n%s", ee, ee.Stderr)
   198  		}
   199  		return nil, err
   200  	}
   201  
   202  	return stdout, nil
   203  }
   204
View as plain text