1+ package org .testcontainers .vault ;
2+
3+ import com .github .dockerjava .api .command .InspectContainerResponse ;
4+ import org .testcontainers .containers .GenericContainer ;
5+ import org .testcontainers .containers .traits .LinkableContainer ;
6+
7+ import java .io .IOException ;
8+ import java .util .ArrayList ;
9+ import java .util .Arrays ;
10+ import java .util .HashMap ;
11+ import java .util .List ;
12+ import java .util .Map ;
13+
14+ import static com .github .dockerjava .api .model .Capability .IPC_LOCK ;
15+
16+
17+ /**
18+ * GenericContainer subclass for Vault specific configuration and features. The main feature is the
19+ * withSecretInVault method, where users can specify which secrets to be pre-loaded into Vault for
20+ * their specific test scenario.
21+ *
22+ * Other helpful features include the withVaultPort, and withVaultToken methods for convenience.
23+ */
24+ public class VaultContainer <SELF extends VaultContainer <SELF >> extends GenericContainer <SELF >
25+ implements LinkableContainer {
26+
27+ private static final String VAULT_PORT = "8200" ;
28+
29+ private boolean vaultPortRequested = false ;
30+
31+ private Map <String , List <String >> secretsMap = new HashMap <>();
32+
33+ public VaultContainer () {
34+ this ("vault:0.7.0" );
35+ }
36+
37+ public VaultContainer (String dockerImageName ) {
38+ super (dockerImageName );
39+ }
40+
41+ @ Override
42+ protected void configure () {
43+ setStartupAttempts (3 );
44+ withCreateContainerCmdModifier (cmd -> cmd .withCapAdd (IPC_LOCK ));
45+ if (!isVaultPortRequested ()){
46+ withEnv ("VAULT_ADDR" , "http://0.0.0.0:" + VAULT_PORT );
47+ }
48+ }
49+
50+ @ Override
51+ protected void containerIsStarted (InspectContainerResponse containerInfo ) {
52+ addSecrets ();
53+ }
54+
55+ private void addSecrets () {
56+ if (!secretsMap .isEmpty ()){
57+ try {
58+ this .execInContainer (buildExecCommand (secretsMap )).getStdout ().contains ("Success" );
59+ }
60+ catch (IOException | InterruptedException e ) {
61+ logger ().error ("Failed to add these secrets {} into Vault via exec command. Exception message: {}" , secretsMap , e .getMessage ());
62+ }
63+ }
64+ }
65+
66+ private String [] buildExecCommand (Map <String , List <String >> map ) {
67+ StringBuilder stringBuilder = new StringBuilder ();
68+ map .forEach ((path , secrets ) -> {
69+ stringBuilder .append (" && vault write " + path );
70+ secrets .forEach (item -> stringBuilder .append (" " + item ));
71+ });
72+ return new String [] { "/bin/sh" , "-c" , stringBuilder .toString ().substring (4 )};
73+ }
74+
75+ /**
76+ * Sets the Vault root token for the container so application tests can source secrets using the token
77+ *
78+ * @param token the root token value to set for Vault.
79+ * @return this
80+ */
81+ public SELF withVaultToken (String token ) {
82+ withEnv ("VAULT_DEV_ROOT_TOKEN_ID" , token );
83+ withEnv ("VAULT_TOKEN" , token );
84+ return self ();
85+ }
86+
87+ /**
88+ * Sets the Vault port in the container as well as the port bindings for the host to reach the container over HTTP.
89+ *
90+ * @param port the port number you want to have the Vault container listen on for tests.
91+ * @return this
92+ */
93+ public SELF withVaultPort (int port ){
94+ setVaultPortRequested (true );
95+ String vaultPort = String .valueOf (port );
96+ withEnv ("VAULT_ADDR" , "http://0.0.0.0:" + VAULT_PORT );
97+ setPortBindings (Arrays .asList (vaultPort + ":" + VAULT_PORT ));
98+ return self ();
99+ }
100+
101+ /**
102+ * Pre-loads secrets into Vault container. User may specify one or more secrets and all will be added to each path
103+ * that is specified. Thus this can be called more than once for multiple paths to be added to Vault.
104+ *
105+ * The secrets are added to vault directly after the container is up via the
106+ * {@link #addSecrets() addSecrets}, called from {@link #containerIsStarted(InspectContainerResponse) containerIsStarted}
107+ *
108+ * @param path specific Vault path to store specified secrets
109+ * @param firstSecret first secret to add to specifed path
110+ * @param remainingSecrets var args list of secrets to add to specified path
111+ * @return this
112+ */
113+ public SELF withSecretInVault (String path , String firstSecret , String ... remainingSecrets ) {
114+ List <String > list = new ArrayList <>();
115+ list .add (firstSecret );
116+ for (String secret : remainingSecrets ) {
117+ list .add (secret );
118+ }
119+ if (secretsMap .containsKey (path )) {
120+ list .addAll (list );
121+ }
122+ secretsMap .putIfAbsent (path ,list );
123+ return self ();
124+ }
125+
126+ private void setVaultPortRequested (boolean vaultPortRequested ) {
127+ this .vaultPortRequested = vaultPortRequested ;
128+ }
129+
130+ private boolean isVaultPortRequested () {
131+ return vaultPortRequested ;
132+ }
133+ }
0 commit comments