diff --git a/plugin_generate.go b/plugin_generate.go
index fe8fef75d..c52958e20 100644
--- a/plugin_generate.go
+++ b/plugin_generate.go
@@ -4,55 +4,16 @@ package main
 
 import (
 	"bytes"
-	"errors"
 	"go/ast"
 	"go/parser"
 	"go/printer"
 	"go/token"
 	"io/ioutil"
 	"log"
-	"strconv"
+
+	"golang.org/x/tools/go/ast/astutil"
 )
 
-func AddImportToFile(file, imprt string) ([]byte, error) {
-	fset := token.NewFileSet()
-	f, err := parser.ParseFile(fset, file, nil, parser.ParseComments)
-	if err != nil {
-		return nil, err
-	}
-
-	for _, s := range f.Imports {
-		iSpec := &ast.ImportSpec{Path: &ast.BasicLit{Value: s.Path.Value}}
-		if iSpec.Path.Value == strconv.Quote(imprt) {
-			return nil, errors.New("coredns import already found")
-		}
-	}
-
-	for i := 0; i < len(f.Decls); i++ {
-		d := f.Decls[i]
-
-		switch d.(type) {
-		case *ast.FuncDecl:
-			// No action
-		case *ast.GenDecl:
-			dd := d.(*ast.GenDecl)
-
-			// IMPORT Declarations
-			if dd.Tok == token.IMPORT {
-				// Add the new import
-				iSpec := &ast.ImportSpec{Name: &ast.Ident{Name: "_"}, Path: &ast.BasicLit{Value: strconv.Quote(imprt)}}
-				dd.Specs = append(dd.Specs, iSpec)
-				break
-			}
-		}
-	}
-
-	ast.SortImports(fset, f)
-
-	out, err := GenerateFile(fset, f)
-	return out, err
-}
-
 func GenerateFile(fset *token.FileSet, file *ast.File) ([]byte, error) {
 	var output []byte
 	buffer := bytes.NewBuffer(output)
@@ -63,19 +24,25 @@ func GenerateFile(fset *token.FileSet, file *ast.File) ([]byte, error) {
 	return buffer.Bytes(), nil
 }
 
-const (
-	coredns = "github.com/miekg/coredns/core"
-	// If everything is OK and we are sitting in CoreDNS' dir, this is where run.go should be.
-	caddyrun = "../../mholt/caddy/caddy/caddymain/run.go"
-)
-
 func main() {
-	out, err := AddImportToFile(caddyrun, coredns)
+	fset := token.NewFileSet()
+	f, err := parser.ParseFile(fset, caddyrun, nil, parser.ParseComments)
 	if err != nil {
-		log.Printf("failed to add import: %s", err)
-		return
+		log.Fatalf("failed to parse %s: %s", caddyrun, err)
 	}
+	astutil.AddNamedImport(fset, f, "_", coredns)
+	astutil.DeleteNamedImport(fset, f, "_", caddy)
+
+	out, err := GenerateFile(fset, f)
 	if err := ioutil.WriteFile(caddyrun, out, 0644); err != nil {
 		log.Fatalf("failed to write go file: %s", err)
 	}
 }
+
+const (
+	coredns = "github.com/miekg/coredns/core"
+	caddy   = "github.com/mholt/caddy/caddyhttp"
+
+	// If everything is OK and we are sitting in CoreDNS' dir, this is where run.go should be.
+	caddyrun = "../../mholt/caddy/caddy/caddymain/run.go"
+)