3333import java .nio .file .Path ;
3434import java .nio .file .Paths ;
3535import java .time .Duration ;
36+ import java .util .Collection ;
3637import java .util .Comparator ;
38+ import java .util .List ;
3739import java .util .Map ;
3840import java .util .Optional ;
3941import java .util .Set ;
5052import org .testcontainers .containers .wait .strategy .WaitAllStrategy ;
5153import org .testcontainers .containers .wait .strategy .WaitStrategy ;
5254import org .testcontainers .containers .wait .strategy .WaitStrategyTarget ;
55+ import org .testcontainers .images .builder .Transferable ;
5356import org .testcontainers .utility .MountableFile ;
5457
5558public class HiveContainer extends BaseContainer {
@@ -70,6 +73,9 @@ public class HiveContainer extends BaseContainer {
7073 public static final int HDFS_DEFAULTFS_PORT = 9000 ;
7174 public static final int HIVE_METASTORE_PORT = 9083 ;
7275 public static final int HIVE_SERVICE_PORT = 10000 ;
76+
77+ /** The realm of the Kerberos image's KDC. */
78+ public static final String KERBEROS_REALM = "HADOOPKRB" ;
7379 // The Kerberos image's second metastore, from its hive-site1.xml.
7480 private static final int SECOND_HIVE_METASTORE_PORT = 19083 ;
7581
@@ -85,8 +91,14 @@ public class HiveContainer extends BaseContainer {
8591 // The image's core-site.xml already lets these users impersonate.
8692 private static final Set <String > USERS_ALLOWED_TO_IMPERSONATE =
8793 ImmutableSet .of ("root" , "hive" , "gravitino" );
94+ // Where the container gets the script that creates a test's Kerberos principals.
95+ static final String PRINCIPALS_SCRIPT = "/tmp/gravitino-kerberos-principals.sh" ;
96+ // The Kerberos image's start script logs in as hive/<host> right after creating its own
97+ // principals and before it configures and starts HDFS; the test's principals go in before it.
98+ static final String PRINCIPALS_HOOK_ANCHOR = "kinit hive/" ;
8899
89100 private final boolean kerberosEnabled ;
101+ private final List <String > kerberosPrincipals ;
90102
91103 public static Builder builder () {
92104 return new Builder ();
@@ -100,9 +112,26 @@ protected HiveContainer(
100112 Map <String , String > filesToMount ,
101113 Map <String , String > envVars ,
102114 Optional <Network > network ,
103- boolean kerberosEnabled ) {
115+ boolean kerberosEnabled ,
116+ List <String > kerberosPrincipals ) {
104117 super (image , hostName , ports , extraHosts , filesToMount , envVars , network );
105118 this .kerberosEnabled = kerberosEnabled ;
119+ this .kerberosPrincipals = List .copyOf (kerberosPrincipals );
120+ if (!kerberosEnabled && !this .kerberosPrincipals .isEmpty ()) {
121+ throw new IllegalArgumentException (
122+ "Kerberos principals need a Kerberos Hive: " + this .kerberosPrincipals );
123+ }
124+ }
125+
126+ /**
127+ * Returns where the Kerberos Hive keeps the keytab of a principal from {@link
128+ * HiveFixture#kerberosPrincipals()}. Tests copy it out with {@code copyFileFromContainer}.
129+ *
130+ * @param principal the principal's name, without realm
131+ * @return the keytab's path in the container
132+ */
133+ public static String kerberosKeytab (String principal ) {
134+ return "/" + principal + ".keytab" ;
106135 }
107136
108137 public String getHostName () {
@@ -140,6 +169,8 @@ protected void setupContainer() {
140169 public void start () {
141170 if (!kerberosEnabled ) {
142171 allowHadoopUserToImpersonate ();
172+ } else if (!kerberosPrincipals .isEmpty ()) {
173+ createKerberosPrincipals ();
143174 }
144175 // Set here rather than in waitStrategy(), which runs before the constructor sets the auth.
145176 container .waitingFor (kerberosEnabled ? kerberosReadiness () : readiness (getHadoopUser ()));
@@ -182,6 +213,62 @@ private void allowHadoopUserToImpersonate() {
182213 cmd -> cmd .withEntrypoint ("/bin/bash" , "-c" , script + " && exec /bin/bash " + IMAGE_START ));
183214 }
184215
216+ /**
217+ * Creates the test's principals and their keytabs while the container starts.
218+ *
219+ * <p>The KDC's database only exists once the image's start script has created it, so the
220+ * entrypoint hooks a script into the start script, next to the image's own principals and before
221+ * HDFS starts; readiness therefore implies the keytabs exist. A missing hook point or a failed
222+ * principal stops the start script, and with it the container.
223+ */
224+ private void createKerberosPrincipals () {
225+ container
226+ .withCopyToContainer (
227+ Transferable .of (principalsScript (kerberosPrincipals ), 0755 ), PRINCIPALS_SCRIPT )
228+ .withCreateContainerCmdModifier (
229+ cmd -> cmd .withEntrypoint ("/bin/bash" , "-c" , principalsEntrypoint (IMAGE_START )));
230+ }
231+
232+ /**
233+ * Returns the script that creates each principal in {@link #KERBEROS_REALM} with a keytab at
234+ * {@link #kerberosKeytab(String)}, and fails if a keytab does not hold its principal. {@code
235+ * kadmin.local -q} exits 0 even when its query fails, so the keytab is the check.
236+ */
237+ static String principalsScript (List <String > principals ) {
238+ StringBuilder script = new StringBuilder ("set -e\n " );
239+ for (String name : principals ) {
240+ String principal = name + "@" + KERBEROS_REALM ;
241+ String keytab = kerberosKeytab (name );
242+ script
243+ .append (format ("kadmin.local -q 'addprinc -randkey %s'\n " , principal ))
244+ .append (format ("kadmin.local -q 'xst -norandkey -k %s %s'\n " , keytab , principal ))
245+ .append (
246+ format (
247+ "klist -k %1$s | grep -q ' %2$s$'"
248+ + " || { echo 'No keytab for %2$s at %1$s' >&2; exit 1; }\n " ,
249+ keytab , principal ));
250+ }
251+ return script .toString ();
252+ }
253+
254+ /**
255+ * Returns the entrypoint that runs {@link #PRINCIPALS_SCRIPT} inside {@code imageStart}, before
256+ * the line with {@link #PRINCIPALS_HOOK_ANCHOR}, and then runs {@code imageStart}. It fails
257+ * without starting anything if that line is missing.
258+ */
259+ static String principalsEntrypoint (String imageStart ) {
260+ String hook =
261+ format (
262+ "bash %s || { echo \" Cannot create the test Kerberos principals\" >&2; exit 1; }" ,
263+ PRINCIPALS_SCRIPT );
264+ return format (
265+ "sed -i '/%1$s/i %2$s' %3$s"
266+ + " && grep -q '^bash %4$s ' %3$s"
267+ + " || { echo 'Cannot hook the Kerberos principals into %3$s' >&2; exit 1; }"
268+ + " && exec /bin/bash %3$s" ,
269+ PRINCIPALS_HOOK_ANCHOR .replace ("/" , "\\ /" ), hook , imageStart , PRINCIPALS_SCRIPT );
270+ }
271+
185272 @ Override
186273 public void close () {
187274 // A private container is closed by its JUnit store and again by the run's closer at the end of
@@ -407,6 +494,9 @@ private static String httpRequest(String method, String url) {
407494 }
408495
409496 public static class Builder extends BaseContainer .Builder <Builder , HiveContainer > {
497+ // Package-private so that tests can read what a spec asked for.
498+ List <String > kerberosPrincipals = List .of ();
499+
410500 private Builder () {
411501 this .image = DEFAULT_IMAGE ;
412502 this .hostName = HOST_NAME ;
@@ -432,7 +522,20 @@ public HiveContainer build() {
432522 filesToMount ,
433523 envVars ,
434524 network ,
435- kerberosEnabled );
525+ kerberosEnabled ,
526+ kerberosPrincipals );
527+ }
528+
529+ /**
530+ * Sets the extra Kerberos principals a Kerberos Hive creates when it starts.
531+ *
532+ * @param kerberosPrincipals the principal names, without realm; see {@link
533+ * HiveFixture#kerberosPrincipals()}
534+ * @return this builder
535+ */
536+ public Builder withKerberosPrincipals (Collection <String > kerberosPrincipals ) {
537+ this .kerberosPrincipals = List .copyOf (kerberosPrincipals );
538+ return this ;
436539 }
437540 }
438541}
0 commit comments