Repository: uber-go/fx
Pull request #824 (2021-12-09 22:01:01)

Validate arguments to fx.Annotate for fx.In and fx.Out
Functions that either take in fx.In as input parameters or outputs
an fx.Out struct should not be used as parameters to fx.Annotate.

This validates the input function to fx.Annotate to check for such
case and bail out early.

Ref: GO-1087

codecov[bot]: # [Codecov](https://codecov.io/gh/uber-go/fx/pull/824?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go) Report
> Merging [#824](https://codecov.io/gh/uber-go/fx/pull/824?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go) (8564961) into [master](https://codecov.io/gh/uber-go/fx/commit/a470642154ec25e847e9243fa9a4b6db799b7b37?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go) (a470642) will **increase** coverage by `0.01%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/uber-go/fx/pull/824/graphs/tree.svg?width=650&height=150&src=pr&token=PYY9PSfaC6&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go)](https://codecov.io/gh/uber-go/fx/pull/824?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go)

```diff
@@            Coverage Diff             @@
##           master     #824      +/-   ##
==========================================
+ Coverage   99.06%   99.07%   +0.01%     
==========================================
  Files          25       25              
  Lines         958      972      +14     
==========================================
+ Hits          949      963      +14     
  Misses          7        7              
  Partials        2        2              
```


| [Impacted Files](https://codecov.io/gh/uber-go/fx/pull/824?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go) | Coverage Δ | |
|---|---|---|
| [annotated.go](https://codecov.io/gh/uber-go/fx/pull/824/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go#diff-YW5ub3RhdGVkLmdv) | `99.38% <100.00%> (+0.05%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/uber-go/fx/pull/824?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/uber-go/fx/pull/824?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go). Last update [a470642...8564961](https://codecov.io/gh/uber-go/fx/pull/824?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=uber-go).


Patch:
diff --git a/annotated.go b/annotated.go
index 2549c80..a3f1aa5 100644
--- a/annotated.go
+++ b/annotated.go
@@ -26,6 +26,7 @@ import (
 	"reflect"
 	"strings"
 
+	"go.uber.org/dig"
 	"go.uber.org/fx/internal/fxreflect"
 )
 
@@ -289,6 +290,10 @@ func (ann *annotated) Build() (interface{}, error) {
 		return nil, fmt.Errorf("must provide constructor function, got %v (%T)", ann.Target, ann.Target)
 	}
 
+	if err := ann.typeCheckOrigFn(); err != nil {
+		return nil, fmt.Errorf("invalid annotation function %T: %w", ann.Target, err)
+	}
+
 	paramTypes, remapParams := ann.parameters()
 	resultTypes, remapResults, err := ann.results()
 	if err != nil {
@@ -312,6 +317,33 @@ func (ann *annotated) Build() (interface{}, error) {
 	return newFn.Interface(), nil
 }
 
+// checks whether the target function is either
+// returning an fx.Out struct or an taking in a
+// fx.In struct as a parameter.
+func (ann *annotated) typeCheckOrigFn() error {
+	ft := reflect.TypeOf(ann.Target)
+	for i := 0; i < ft.NumOut(); i++ {
+		ot := ft.Out(i)
+		if ot.Kind() != reflect.Struct {
+			continue
+		}
+		if dig.IsOut(reflect.New(ft.Out(i)).Elem().Interface()) {
+			return errors.New("fx.Out structs cannot be annotated")
+		}
+	}
+
+	for i := 0; i < ft.NumIn(); i++ {
+		it := ft.In(i)
+		if it.Kind() != reflect.Struct {
+			continue
+		}
+		if dig.IsIn(reflect.New(ft.In(i)).Elem().Interface()) {
+			return errors.New("fx.In structs cannot be annotated")
+		}
+	}
+	return nil
+}
+
 // parameters returns the type for the parameters of the annotated function,
 // and a function that maps the arguments of the annotated function
 // back to the arguments of the target function.


Test patch:
diff --git a/annotated_test.go b/annotated_test.go
index 8231c60..583c890 100644
--- a/annotated_test.go
+++ b/annotated_test.go
@@ -850,4 +850,51 @@ func TestAnnotate(t *testing.T) {
 		require.Error(t, err)
 		assert.Contains(t, err.Error(), "must provide constructor function, got 42 (int)")
 	})
+
+	t.Run("annotate a fx.Out", func(t *testing.T) {
+		t.Parallel()
+
+		type A struct {
+			s string
+
+			fx.Out
+		}
+
+		f := func() A {
+			return A{s: "hi"}
+		}
+
+		app := NewForTest(t,
+			fx.Provide(
+				fx.Annotate(f, fx.ResultTags(`name:"out"`)),
+			),
+		)
+
+		err := app.Err()
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "fx.Out structs cannot be annotated")
+	})
+
+	t.Run("annotate a fx.In", func(t *testing.T) {
+		t.Parallel()
+
+		type A struct {
+			S string
+		}
+		type B struct {
+			fx.In
+		}
+
+		app := NewForTest(t,
+			fx.Provide(
+				fx.Annotate(func(i A) string { return i.S }, fx.ParamTags(`optional:"true"`)),
+				fx.Annotate(func(i B) string { return "ok" }, fx.ParamTags(`name:"problem"`)),
+			),
+		)
+		err := app.Err()
+		require.Error(t, err)
+		assert.NotContains(t, err.Error(), "invalid annotation function func(fx_test.A) string")
+		assert.Contains(t, err.Error(), "invalid annotation function func(fx_test.B) string")
+		assert.Contains(t, err.Error(), "fx.In structs cannot be annotated")
+	})
 }